Closes #20129: Enable dynamic model feature registration (#20130)

* Closes #20129: Enable dynamic model feature registration

* Correct import path for register_model_feature()

netbox/core/models/object_types.py

@@ -135,9 +135,9 @@ class ObjectTypeManager(models.Manager):
         """
         Return ObjectTypes only for models which support the given feature.
 
-        Only ObjectTypes which list the specified feature will be included. Supported features are declared in
-        netbox.models.features.FEATURES_MAP. For example, we can find all ObjectTypes for models which support event
-        rules with:
+        Only ObjectTypes which list the specified feature will be included. Supported features are declared in the
+        application registry under `registry["model_features"]`. For example, we can find all ObjectTypes for models
+        which support event rules with:
 
             ObjectType.objects.with_feature('event_rules')
         """

netbox/netbox/models/features.py

@@ -22,6 +22,7 @@ from netbox.models.deletion import DeleteMixin
 from netbox.plugins import PluginConfig
 from netbox.registry import registry
 from netbox.signals import post_clean
+from netbox.utils import register_model_feature
 from utilities.json import CustomFieldJSONEncoder
 from utilities.serialization import serialize_object
 
@@ -35,7 +36,6 @@ __all__ = (
     'CustomValidationMixin',
     'EventRulesMixin',
     'ExportTemplatesMixin',
-    'FEATURES_MAP',
     'ImageAttachmentsMixin',
     'JobsMixin',
     'JournalingMixin',
@@ -628,28 +628,21 @@ class SyncedDataMixin(models.Model):
 # Feature registration
 #
 
-FEATURES_MAP = {
-    'bookmarks': BookmarksMixin,
-    'change_logging': ChangeLoggingMixin,
-    'cloning': CloningMixin,
-    'contacts': ContactsMixin,
-    'custom_fields': CustomFieldsMixin,
-    'custom_links': CustomLinksMixin,
-    'custom_validation': CustomValidationMixin,
-    'event_rules': EventRulesMixin,
-    'export_templates': ExportTemplatesMixin,
-    'image_attachments': ImageAttachmentsMixin,
-    'jobs': JobsMixin,
-    'journaling': JournalingMixin,
-    'notifications': NotificationsMixin,
-    'synced_data': SyncedDataMixin,
-    'tags': TagsMixin,
-}
-
-# TODO: Remove in NetBox v4.5
-registry['model_features'].update({
-    feature: defaultdict(set) for feature in FEATURES_MAP.keys()
-})
+register_model_feature('bookmarks', lambda model: issubclass(model, BookmarksMixin))
+register_model_feature('change_logging', lambda model: issubclass(model, ChangeLoggingMixin))
+register_model_feature('cloning', lambda model: issubclass(model, CloningMixin))
+register_model_feature('contacts', lambda model: issubclass(model, ContactsMixin))
+register_model_feature('custom_fields', lambda model: issubclass(model, CustomFieldsMixin))
+register_model_feature('custom_links', lambda model: issubclass(model, CustomLinksMixin))
+register_model_feature('custom_validation', lambda model: issubclass(model, CustomValidationMixin))
+register_model_feature('event_rules', lambda model: issubclass(model, EventRulesMixin))
+register_model_feature('export_templates', lambda model: issubclass(model, ExportTemplatesMixin))
+register_model_feature('image_attachments', lambda model: issubclass(model, ImageAttachmentsMixin))
+register_model_feature('jobs', lambda model: issubclass(model, JobsMixin))
+register_model_feature('journaling', lambda model: issubclass(model, JournalingMixin))
+register_model_feature('notifications', lambda model: issubclass(model, NotificationsMixin))
+register_model_feature('synced_data', lambda model: issubclass(model, SyncedDataMixin))
+register_model_feature('tags', lambda model: issubclass(model, TagsMixin))
 
 
 def model_is_public(model):
@@ -665,8 +658,11 @@ def model_is_public(model):
 
 
 def get_model_features(model):
+    """
+    Return all features supported by the given model.
+    """
     return [
-        feature for feature, cls in FEATURES_MAP.items() if issubclass(model, cls)
+        feature for feature, test_func in registry['model_features'].items() if test_func(model)
     ]
 
 
@@ -710,19 +706,6 @@ def register_models(*models):
         if not getattr(model, '_netbox_private', False):
             registry['models'][app_label].add(model_name)
 
-        # TODO: Remove in NetBox v4.5
-        # Record each applicable feature for the model in the registry
-        features = {
-            feature for feature, cls in FEATURES_MAP.items() if issubclass(model, cls)
-        }
-        for feature in features:
-            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()}"
-                )
-
         # Register applicable feature views for the model
         if issubclass(model, ContactsMixin):
             register_model_view(model, 'contacts', kwargs={'model': model})(

netbox/netbox/utils.py

@@ -3,6 +3,7 @@ from netbox.registry import registry
 __all__ = (
     'get_data_backend_choices',
     'register_data_backend',
+    'register_model_feature',
     'register_request_processor',
 )
 
@@ -27,6 +28,35 @@ def register_data_backend():
     return _wrapper
 
 
+def register_model_feature(name, func=None):
+    """
+    Register a model feature with its qualifying function.
+
+    The qualifying function must accept a single `model` argument. It will be called to determine whether the given
+    model supports the corresponding feature.
+
+    This function can be used directly:
+
+        register_model_feature('my_feature', my_func)
+
+    Or as a decorator:
+
+        @register_model_feature('my_feature')
+        def my_func(model):
+            ...
+    """
+    def decorator(f):
+        registry['model_features'][name] = f
+        return f
+
+    if name in registry['model_features']:
+        raise ValueError(f"A model feature named {name} is already registered.")
+
+    if func is None:
+        return decorator
+    return decorator(func)
+
+
 def register_request_processor(func):
     """
     Decorator for registering a request processor.