Merge pull request #4379 from netbox-community/refactor-registry

Refactor registry to behave as a dictionary

netbox/extras/registry.py

@@ -0,0 +1,21 @@
+class Registry(dict):
+    """
+    Central registry for registration of functionality. Once a store (key) is defined, it cannot be overwritten or
+    deleted (although its value may be manipulated).
+    """
+    def __getitem__(self, key):
+        try:
+            return super().__getitem__(key)
+        except KeyError:
+            raise KeyError("Invalid store: {}".format(key))
+
+    def __setitem__(self, key, value):
+        if key in self:
+            raise KeyError("Store already set: {}".format(key))
+        super().__setitem__(key, value)
+
+    def __delitem__(self, key):
+        raise TypeError("Cannot delete stores from registry")
+
+
+registry = Registry()

netbox/extras/tests/test_registry.py

@@ -0,0 +1,33 @@
+from django.test import TestCase
+
+from extras.registry import Registry
+
+
+class RegistryTest(TestCase):
+
+    def test_add_store(self):
+        reg = Registry()
+        reg['foo'] = 123
+
+        self.assertEqual(reg['foo'], 123)
+
+    def test_manipulate_store(self):
+        reg = Registry()
+        reg['foo'] = [1, 2]
+        reg['foo'].append(3)
+
+        self.assertListEqual(reg['foo'], [1, 2, 3])
+
+    def test_overwrite_store(self):
+        reg = Registry()
+        reg['foo'] = 123
+
+        with self.assertRaises(KeyError):
+            reg['foo'] = 456
+
+    def test_delete_store(self):
+        reg = Registry()
+        reg['foo'] = 123
+
+        with self.assertRaises(TypeError):
+            del(reg['foo'])

netbox/extras/utils.py

@@ -6,6 +6,7 @@ from taggit.managers import _TaggableManager
 from utilities.querysets import DummyQuerySet
 
 from extras.constants import EXTRAS_FEATURES
+from extras.registry import registry
 
 
 def is_taggable(obj):
@@ -21,33 +22,12 @@ def is_taggable(obj):
     return False
 
 
-#
-# Dynamic feature registration
-#
-
-class Registry:
-    """
-    The registry is a place to hook into for data storage across components
-    """
-
-    def add_store(self, store_name, initial_value=None):
-        """
-        Given the name of some new data parameter and an optional initial value, setup the registry store
-        """
-        if not hasattr(Registry, store_name):
-            setattr(Registry, store_name, initial_value)
-
-
-registry = Registry()
-
-
 @deconstructible
 class FeatureQuery:
     """
-    Helper class that delays evaluation of the registry contents for the functionaility store
+    Helper class that delays evaluation of the registry contents for the functionality store
     until it has been populated.
     """
-
     def __init__(self, feature):
         self.feature = feature
 
@@ -59,24 +39,26 @@ class FeatureQuery:
         Given an extras feature, return a Q object for content type lookup
         """
         query = Q()
-        for app_label, models in registry.model_feature_store[self.feature].items():
+        for app_label, models in registry['model_features'][self.feature].items():
             query |= Q(app_label=app_label, model__in=models)
 
         return query
 
 
-registry.add_store('model_feature_store', {f: collections.defaultdict(list) for f in EXTRAS_FEATURES})
-
-
 def extras_features(*features):
     """
     Decorator used to register extras provided features to a model
     """
     def wrapper(model_class):
+        # Initialize the model_features store if not already defined
+        if 'model_features' not in registry:
+            registry['model_features'] = {
+                f: collections.defaultdict(list) for f in EXTRAS_FEATURES
+            }
         for feature in features:
             if feature in EXTRAS_FEATURES:
                 app_label, model_name = model_class._meta.label_lower.split('.')
-                registry.model_feature_store[feature][app_label].append(model_name)
+                registry['model_features'][feature][app_label].append(model_name)
             else:
                 raise ValueError('{} is not a valid extras feature!'.format(feature))
         return model_class