Clean up the application registry

Author: jeremystretch

netbox/extras/constants.py

@@ -1,14 +1,2 @@
 # Webhook content types
 HTTP_CONTENT_TYPE_JSON = 'application/json'
-
-# Registerable extras features
-EXTRAS_FEATURES = [
-    'custom_fields',
-    'custom_links',
-    'export_templates',
-    'job_results',
-    'journaling',
-    'synced_data',
-    'tags',
-    'webhooks'
-]

netbox/extras/plugins/__init__.py

@@ -14,13 +14,13 @@
 from .templates import *
 
 # Initialize plugin registry
-registry['plugins'] = {
+registry['plugins'].update({
     'graphql_schemas': [],
     'menus': [],
     'menu_items': {},
     'preferences': {},
     'template_extensions': collections.defaultdict(list),
-}
+})
 
 DEFAULT_RESOURCE_PATHS = {
     'search_indexes': 'search.indexes',

netbox/extras/utils.py

@@ -2,7 +2,6 @@
 from django.utils.deconstruct import deconstructible
 from taggit.managers import _TaggableManager
 
-from extras.constants import EXTRAS_FEATURES
 from netbox.registry import registry
 
 
@@ -18,7 +17,7 @@ def is_taggable(obj):
 
 def image_upload(instance, filename):
     """
-    Return a path for uploading image attchments.
+    Return a path for uploading image attachments.
     """
     path = 'image-attachments/'
 
@@ -56,8 +55,14 @@ def get_query(self):
 
 
 def register_features(model, features):
+    """
+    Register model features in the application registry.
+    """
+    app_label, model_name = model._meta.label_lower.split('.')
     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].add(model_name)
+        try:
+            registry['model_features'][feature][app_label].add(model_name)
+        except KeyError:
+            raise KeyError(
+                f"{feature} is not a valid model feature! Valid keys are: {registry['model_features'].keys()}"
+            )

netbox/netbox/models/features.py

@@ -12,6 +12,7 @@
 
 from extras.choices import CustomFieldVisibilityChoices, ObjectChangeActionChoices
 from extras.utils import is_taggable, register_features
+from netbox.registry import registry
 from netbox.signals import post_clean
 from utilities.json import CustomFieldJSONEncoder
 from utilities.utils import serialize_object
@@ -388,22 +389,26 @@ def sync_data(self):
         raise NotImplementedError(f"{self.__class__} must implement a sync_data() method.")
 
 
-FEATURES_MAP = (
-    ('custom_fields', CustomFieldsMixin),
-    ('custom_links', CustomLinksMixin),
-    ('export_templates', ExportTemplatesMixin),
-    ('job_results', JobResultsMixin),
-    ('journaling', JournalingMixin),
-    ('synced_data', SyncedDataMixin),
-    ('tags', TagsMixin),
-    ('webhooks', WebhooksMixin),
-)
+FEATURES_MAP = {
+    'custom_fields': CustomFieldsMixin,
+    'custom_links': CustomLinksMixin,
+    'export_templates': ExportTemplatesMixin,
+    'job_results': JobResultsMixin,
+    'journaling': JournalingMixin,
+    'synced_data': SyncedDataMixin,
+    'tags': TagsMixin,
+    'webhooks': WebhooksMixin,
+}
+
+registry['model_features'].update({
+    feature: defaultdict(set) for feature in FEATURES_MAP.keys()
+})
 
 
 @receiver(class_prepared)
 def _register_features(sender, **kwargs):
     features = {
-        feature for feature, cls in FEATURES_MAP if issubclass(sender, cls)
+        feature for feature, cls in FEATURES_MAP.items() if issubclass(sender, cls)
     }
     register_features(sender, features)
 

netbox/netbox/registry.py

@@ -1,12 +1,10 @@
 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
-    deleted (although its value may be manipulated).
+    Central registry for registration of functionality. Once a Registry is initialized, keys cannot be added or
+    removed (though the value of each key is mutable).
     """
     def __getitem__(self, key):
         try:
@@ -15,20 +13,18 @@ def __getitem__(self, key):
             raise KeyError(f"Invalid store: {key}")
 
     def __setitem__(self, key, value):
-        if key in self:
-            raise KeyError(f"Store already set: {key}")
-        super().__setitem__(key, value)
+        raise TypeError("Cannot add stores to registry after initialization")
 
     def __delitem__(self, key):
         raise TypeError("Cannot delete stores from registry")
 
 
 # Initialize the global registry
-registry = Registry()
-registry['data_backends'] = dict()
-registry['denormalized_fields'] = collections.defaultdict(list)
-registry['model_features'] = {
-    feature: collections.defaultdict(set) for feature in EXTRAS_FEATURES
-}
-registry['search'] = dict()
-registry['views'] = collections.defaultdict(dict)
+registry = Registry({
+    'data_backends': dict(),
+    'denormalized_fields': collections.defaultdict(list),
+    'model_features': dict(),
+    'plugins': dict(),
+    'search': dict(),
+    'views': collections.defaultdict(dict),
+})

netbox/netbox/tests/test_registry.py

@@ -5,29 +5,23 @@
 
 class RegistryTest(TestCase):
 
-    def test_add_store(self):
-        reg = Registry()
-        reg['foo'] = 123
-
-        self.assertEqual(reg['foo'], 123)
+    def test_set_store(self):
+        reg = Registry({
+            'foo': 123,
+        })
+        with self.assertRaises(TypeError):
+            reg['bar'] = 456
 
-    def test_manipulate_store(self):
-        reg = Registry()
-        reg['foo'] = [1, 2]
+    def test_mutate_store(self):
+        reg = Registry({
+            '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
-
+        reg = Registry({
+            'foo': 123,
+        })
         with self.assertRaises(TypeError):
             del reg['foo']