diff --git a/netbox/core/management/commands/nbshell.py b/netbox/core/management/commands/nbshell.py index 0eb1b0564..833803b67 100644 --- a/netbox/core/management/commands/nbshell.py +++ b/netbox/core/management/commands/nbshell.py @@ -1,18 +1,18 @@ import code import platform -import sys +from collections import defaultdict +from types import SimpleNamespace from colorama import Fore, Style from django import get_version from django.apps import apps from django.conf import settings from django.core.management.base import BaseCommand +from django.utils.module_loading import import_string from netbox.constants import CORE_APPS from netbox.plugins.utils import get_installed_plugins -EXCLUDE_MODELS = () - def color(color: str, text: str): return getattr(Fore, color.upper()) + text + Style.RESET_ALL @@ -22,6 +22,29 @@ def bright(text: str): return Style.BRIGHT + text + Style.RESET_ALL +def get_models(app_config): + """ + Return a list of all non-private models within an app. + """ + return [ + model for model in app_config.get_models() + if not getattr(model, '_netbox_private', False) + ] + + +def get_constants(app_config): + """ + Return a dictionary mapping of all constants defined within an app. + """ + try: + constants = import_string(f'{app_config.name}.constants') + except ImportError: + return {} + return { + name: value for name, value in vars(constants).items() + } + + class Command(BaseCommand): help = "Start the Django shell with all NetBox models already imported" django_models = {} @@ -40,37 +63,31 @@ class Command(BaseCommand): print(f' {m}') def get_namespace(self): - namespace = {} + namespace = defaultdict(SimpleNamespace) - # Gather Django models and constants from each app - for app in CORE_APPS: - models = [] + # Iterate through all core apps & plugins to compile namespace of models and constants + for app_name in [*CORE_APPS, *get_installed_plugins().keys()]: + app_config = apps.get_app_config(app_name) - # Load models from each app - for model in apps.get_app_config(app).get_models(): - app_label = model._meta.app_label - model_name = model._meta.model_name - if f'{app_label}.{model_name}' not in EXCLUDE_MODELS: - namespace[model.__name__] = model - models.append(model.__name__) - self.django_models[app] = sorted(models) + # Populate models + if models := get_models(app_config): + for model in models: + setattr(namespace[app_name], model.__name__, model) + self.django_models[app_name] = sorted([ + model.__name__ for model in models + ]) - # Constants - try: - app_constants = sys.modules[f'{app}.constants'] - for name in dir(app_constants): - namespace[name] = getattr(app_constants, name) - except KeyError: - pass + # Populate constants + for const_name, const_value in get_constants(app_config).items(): + setattr(namespace[app_name], const_name, const_value) - # Load convenience commands - namespace.update({ + return { + **namespace, 'lsmodels': self._lsmodels, - }) + } - return namespace - - def get_banner_text(self): + @staticmethod + def get_banner_text(): lines = [ '{title} ({hostname})'.format( title=bright('NetBox interactive shell'), @@ -120,5 +137,4 @@ class Command(BaseCommand): readline.parse_and_bind('tab: complete') # Run interactive shell - shell = code.interact(banner=self.get_banner_text(), local=namespace) - return shell + return code.interact(banner=self.get_banner_text(), local=namespace)