diff --git a/netbox/core/data_backends.py b/netbox/core/data_backends.py index c0063a34f..5d0e80584 100644 --- a/netbox/core/data_backends.py +++ b/netbox/core/data_backends.py @@ -8,6 +8,8 @@ from django import forms from django.conf import settings from django.utils.translation import gettext as _ +from netbox.registry import registry +from .choices import DataSourceTypeChoices from .exceptions import SyncError __all__ = ( @@ -18,6 +20,17 @@ __all__ = ( logger = logging.getLogger('netbox.data_backends') +def register_backend(name): + """ + Decorator for registering a DataBackend class. + """ + def _wrapper(cls): + registry['data_backends'][name] = cls + return cls + + return _wrapper + + class DataBackend: parameters = {} @@ -34,6 +47,7 @@ class DataBackend: raise NotImplemented() +@register_backend(DataSourceTypeChoices.LOCAL) class LocalBackend(DataBackend): @contextmanager @@ -44,6 +58,7 @@ class LocalBackend(DataBackend): yield local_path +@register_backend(DataSourceTypeChoices.GIT) class GitBackend(DataBackend): parameters = { 'username': forms.CharField( diff --git a/netbox/core/forms/model_forms.py b/netbox/core/forms/model_forms.py index b84cbf932..cc2f6c0a4 100644 --- a/netbox/core/forms/model_forms.py +++ b/netbox/core/forms/model_forms.py @@ -2,7 +2,7 @@ import copy from core.models import * from netbox.forms import NetBoxModelForm, StaticSelect -from ..models.data import BACKEND_CLASSES +from netbox.registry import registry __all__ = ( 'DataSourceForm', @@ -33,13 +33,14 @@ class DataSourceForm(NetBoxModelForm): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - if self.is_bound and self.data.get('type') in BACKEND_CLASSES: - backend_type = self.data['type'] - elif self.initial and self.initial.get('type') in BACKEND_CLASSES: - backend_type = self.initial['type'] + backend_classes = registry['data_backends'] + + if self.is_bound and self.data.get('type') in backend_classes: + backend = backend_classes.get(self.data['type']) + elif self.initial and self.initial.get('type') in backend_classes: + backend = backend_classes.get(self.initial['type']) else: - backend_type = self.fields['type'].initial - backend = BACKEND_CLASSES.get(backend_type) + backend = backend_classes.get(self.fields['type'].initial) for name, form_field in backend.parameters.items(): field_name = f'backend_{name}' self.fields[field_name] = copy.copy(form_field) diff --git a/netbox/core/models/data.py b/netbox/core/models/data.py index d3a70f67a..c190d154f 100644 --- a/netbox/core/models/data.py +++ b/netbox/core/models/data.py @@ -14,10 +14,10 @@ from django.utils.translation import gettext as _ from extras.models import JobResult from netbox.models import ChangeLoggedModel +from netbox.registry import registry from utilities.files import sha256_hash from utilities.querysets import RestrictedQuerySet from ..choices import * -from ..data_backends import GitBackend, LocalBackend from ..exceptions import SyncError __all__ = ( @@ -27,11 +27,6 @@ __all__ = ( logger = logging.getLogger('netbox.core.data') -BACKEND_CLASSES = { - DataSourceTypeChoices.LOCAL: LocalBackend, - DataSourceTypeChoices.GIT: GitBackend, -} - class DataSource(ChangeLoggedModel): """ @@ -129,7 +124,7 @@ class DataSource(ChangeLoggedModel): return job_result def get_backend(self): - backend_cls = BACKEND_CLASSES.get(self.type) + backend_cls = registry['data_backends'].get(self.type) backend_params = self.parameters or {} return backend_cls(self.url, **backend_params) diff --git a/netbox/netbox/registry.py b/netbox/netbox/registry.py index 76886e791..670bca683 100644 --- a/netbox/netbox/registry.py +++ b/netbox/netbox/registry.py @@ -25,9 +25,10 @@ class Registry(dict): # 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['denormalized_fields'] = collections.defaultdict(list) registry['search'] = dict() registry['views'] = collections.defaultdict(dict)