diff --git a/netbox/extras/registry.py b/netbox/extras/registry.py index cb58f5135..07fd4cc24 100644 --- a/netbox/extras/registry.py +++ b/netbox/extras/registry.py @@ -1,3 +1,8 @@ +import collections + +from extras.constants import EXTRAS_FEATURES + + class Registry(dict): """ Central registry for registration of functionality. Once a store (key) is defined, it cannot be overwritten or @@ -7,15 +12,19 @@ class Registry(dict): try: return super().__getitem__(key) except KeyError: - raise KeyError("Invalid store: {}".format(key)) + raise KeyError(f"Invalid store: {key}") def __setitem__(self, key, value): if key in self: - raise KeyError("Store already set: {}".format(key)) + raise KeyError(f"Store already set: {key}") super().__setitem__(key, value) def __delitem__(self, key): raise TypeError("Cannot delete stores from registry") +# Initialize the global registry registry = Registry() +registry['model_features'] = { + feature: collections.defaultdict(set) for feature in EXTRAS_FEATURES +} diff --git a/netbox/extras/utils.py b/netbox/extras/utils.py index 487ca3c0b..e16807821 100644 --- a/netbox/extras/utils.py +++ b/netbox/extras/utils.py @@ -1,5 +1,3 @@ -import collections - from django.db.models import Q from django.utils.deconstruct import deconstructible from taggit.managers import _TaggableManager @@ -58,12 +56,8 @@ class FeatureQuery: def register_features(model, features): - if 'model_features' not in registry: - registry['model_features'] = { - f: collections.defaultdict(list) for f in EXTRAS_FEATURES - } for feature in features: if feature not in EXTRAS_FEATURES: raise ValueError(f"{feature} is not a valid extras feature!") app_label, model_name = model._meta.label_lower.split('.') - registry['model_features'][feature][app_label].append(model_name) + registry['model_features'][feature][app_label].add(model_name)