Selaa lähdekoodia

Tweak registry initialization

jeremystretch 4 vuotta sitten
vanhempi
commit
047bed2a86
2 muutettua tiedostoa jossa 12 lisäystä ja 9 poistoa
  1. 11 2
      netbox/extras/registry.py
  2. 1 7
      netbox/extras/utils.py

+ 11 - 2
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
+}

+ 1 - 7
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)