From 3b10cf153cbb4048efed430411333e2b2ce1e7e0 Mon Sep 17 00:00:00 2001 From: Jeremy Stretch Date: Wed, 21 Feb 2024 12:07:20 -0500 Subject: [PATCH] Refactor register_model() to accept multiple models --- netbox/circuits/apps.py | 5 +-- netbox/core/apps.py | 5 +-- netbox/dcim/apps.py | 5 +-- netbox/extras/apps.py | 5 +-- netbox/ipam/apps.py | 5 +-- netbox/netbox/models/features.py | 71 ++++++++++++++++--------------- netbox/netbox/plugins/__init__.py | 5 +-- netbox/tenancy/apps.py | 5 +-- netbox/users/apps.py | 5 +-- netbox/virtualization/apps.py | 5 +-- netbox/vpn/apps.py | 5 +-- netbox/wireless/apps.py | 5 +-- 12 files changed, 58 insertions(+), 68 deletions(-) diff --git a/netbox/circuits/apps.py b/netbox/circuits/apps.py index 176f3b8f1..df6804303 100644 --- a/netbox/circuits/apps.py +++ b/netbox/circuits/apps.py @@ -6,9 +6,8 @@ class CircuitsConfig(AppConfig): verbose_name = "Circuits" def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from . import signals, search # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) diff --git a/netbox/core/apps.py b/netbox/core/apps.py index a198f3856..b1103469c 100644 --- a/netbox/core/apps.py +++ b/netbox/core/apps.py @@ -17,9 +17,8 @@ class CoreConfig(AppConfig): def ready(self): from core.api import schema # noqa - from netbox.models.features import register_model + from netbox.models.features import register_models from . import data_backends, search # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) diff --git a/netbox/dcim/apps.py b/netbox/dcim/apps.py index 631ea00da..4df66e367 100644 --- a/netbox/dcim/apps.py +++ b/netbox/dcim/apps.py @@ -8,14 +8,13 @@ class DCIMConfig(AppConfig): verbose_name = "DCIM" def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from utilities.counters import connect_counters from . import signals, search from .models import CableTermination, Device, DeviceType, VirtualChassis # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) # Register denormalized fields denormalized.register(CableTermination, '_device', { diff --git a/netbox/extras/apps.py b/netbox/extras/apps.py index 40c00fb15..c565988bc 100644 --- a/netbox/extras/apps.py +++ b/netbox/extras/apps.py @@ -5,9 +5,8 @@ class ExtrasConfig(AppConfig): name = "extras" def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from . import dashboard, lookups, search, signals # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) diff --git a/netbox/ipam/apps.py b/netbox/ipam/apps.py index 3f467758e..244ec7d6d 100644 --- a/netbox/ipam/apps.py +++ b/netbox/ipam/apps.py @@ -6,9 +6,8 @@ class IPAMConfig(AppConfig): verbose_name = "IPAM" def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from . import signals, search # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) diff --git a/netbox/netbox/models/features.py b/netbox/netbox/models/features.py index ad6ab74b5..91765be04 100644 --- a/netbox/netbox/models/features.py +++ b/netbox/netbox/models/features.py @@ -35,7 +35,7 @@ __all__ = ( 'JournalingMixin', 'SyncedDataMixin', 'TagsMixin', - 'register_model', + 'register_models', ) @@ -575,9 +575,9 @@ registry['model_features'].update({ }) -def register_model(model, **kwargs): +def register_models(*models): """ - Register a model in NetBox. This entails: + Register one or more models in NetBox. This entails: - Determining whether the model is considered "public" (available for reference by other models) - Registering which features the model supports (e.g. bookmarks, custom fields, etc.) @@ -585,38 +585,39 @@ def register_model(model, **kwargs): register_model() should be called for each relevant model under the ready() of an app's AppConfig class. """ - app_label, model_name = model._meta.label_lower.split('.') + for model in models: + app_label, model_name = model._meta.label_lower.split('.') - # Register public models - if not getattr(model, '_netbox_private', False): - registry['models'][app_label].add(model_name) + # Register public models + if not getattr(model, '_netbox_private', False): + registry['models'][app_label].add(model_name) - # 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()}" + # 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, JournalingMixin): + register_model_view(model, 'journal', kwargs={'model': model})( + 'netbox.views.generic.ObjectJournalView' + ) + if issubclass(model, ChangeLoggingMixin): + register_model_view(model, 'changelog', kwargs={'model': model})( + 'netbox.views.generic.ObjectChangeLogView' + ) + if issubclass(model, JobsMixin): + register_model_view(model, 'jobs', kwargs={'model': model})( + 'netbox.views.generic.ObjectJobsView' + ) + if issubclass(model, SyncedDataMixin): + register_model_view(model, 'sync', kwargs={'model': model})( + 'netbox.views.generic.ObjectSyncDataView' ) - - # Register applicable feature views for the model - if issubclass(model, JournalingMixin): - register_model_view(model, 'journal', kwargs={'model': model})( - 'netbox.views.generic.ObjectJournalView' - ) - if issubclass(model, ChangeLoggingMixin): - register_model_view(model, 'changelog', kwargs={'model': model})( - 'netbox.views.generic.ObjectChangeLogView' - ) - if issubclass(model, JobsMixin): - register_model_view(model, 'jobs', kwargs={'model': model})( - 'netbox.views.generic.ObjectJobsView' - ) - if issubclass(model, SyncedDataMixin): - register_model_view(model, 'sync', kwargs={'model': model})( - 'netbox.views.generic.ObjectSyncDataView' - ) diff --git a/netbox/netbox/plugins/__init__.py b/netbox/netbox/plugins/__init__.py index df210c63a..8db945e74 100644 --- a/netbox/netbox/plugins/__init__.py +++ b/netbox/netbox/plugins/__init__.py @@ -94,11 +94,10 @@ class PluginConfig(AppConfig): pass def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) plugin_name = self.name.rsplit('.', 1)[-1] diff --git a/netbox/tenancy/apps.py b/netbox/tenancy/apps.py index 9e0e6d306..7fec97d73 100644 --- a/netbox/tenancy/apps.py +++ b/netbox/tenancy/apps.py @@ -5,9 +5,8 @@ class TenancyConfig(AppConfig): name = 'tenancy' def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from . import search # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) diff --git a/netbox/users/apps.py b/netbox/users/apps.py index 90552b49c..cd61a59bc 100644 --- a/netbox/users/apps.py +++ b/netbox/users/apps.py @@ -5,9 +5,8 @@ class UsersConfig(AppConfig): name = 'users' def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from . import signals # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) diff --git a/netbox/virtualization/apps.py b/netbox/virtualization/apps.py index 0608d4504..8a61bc523 100644 --- a/netbox/virtualization/apps.py +++ b/netbox/virtualization/apps.py @@ -7,14 +7,13 @@ class VirtualizationConfig(AppConfig): name = 'virtualization' def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from utilities.counters import connect_counters from . import search, signals from .models import VirtualMachine # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) # Register denormalized fields denormalized.register(VirtualMachine, 'cluster', { diff --git a/netbox/vpn/apps.py b/netbox/vpn/apps.py index 094d2e90b..9a3751cb2 100644 --- a/netbox/vpn/apps.py +++ b/netbox/vpn/apps.py @@ -6,9 +6,8 @@ class VPNConfig(AppConfig): verbose_name = 'VPN' def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from . import search # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models()) diff --git a/netbox/wireless/apps.py b/netbox/wireless/apps.py index 9861f49f2..24e1869b6 100644 --- a/netbox/wireless/apps.py +++ b/netbox/wireless/apps.py @@ -5,9 +5,8 @@ class WirelessConfig(AppConfig): name = 'wireless' def ready(self): - from netbox.models.features import register_model + from netbox.models.features import register_models from . import signals, search # Register models - for model in self.get_models(): - register_model(model) + register_models(*self.get_models())