Closes #19973: nbshell improvements (#19995)

base_requirements.txt

@@ -1,3 +1,7 @@
+# Shell text coloring
+# https://github.com/tartley/colorama/blob/master/CHANGELOG.rst
+colorama
+
 # The Python web framework on which NetBox is built
 # https://docs.djangoproject.com/en/stable/releases/
 Django==5.2.*

netbox/core/management/commands/nbshell.py

@@ -1,29 +1,48 @@
 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 core.models import ObjectType
-from users.models import User
+from netbox.constants import CORE_APPS
+from netbox.plugins.utils import get_installed_plugins
 
-APPS = ('circuits', 'core', 'dcim', 'extras', 'ipam', 'tenancy', 'users', 'virtualization', 'vpn', 'wireless')
-EXCLUDE_MODELS = (
-    'extras.branch',
-    'extras.stagedchange',
-)
 
-BANNER_TEXT = """### NetBox interactive shell ({node})
-### Python {python} | Django {django} | NetBox {netbox}
-### lsmodels() will show available models. Use help(<model>) for more info.""".format(
-    node=platform.node(),
-    python=platform.python_version(),
-    django=get_version(),
-    netbox=settings.RELEASE.name
-)
+def color(color: str, text: str):
+    return getattr(Fore, color.upper()) + text + Style.RESET_ALL
+
+
+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):
@@ -36,47 +55,88 @@ class Command(BaseCommand):
             help='Python code to execute (instead of starting an interactive shell)',
         )
 
-    def _lsmodels(self):
-        for app, models in self.django_models.items():
-            app_name = apps.get_app_config(app).verbose_name
+    def _lsapps(self):
+        for app_label in self.django_models.keys():
+            app_name = apps.get_app_config(app_label).verbose_name
+            print(f'{app_label} - {app_name}')
+
+    def _lsmodels(self, app_label=None):
+        """
+        Return a list of all models within each app.
+
+        Args:
+            app_label: The name of a specific app
+        """
+        if app_label:
+            if app_label not in self.django_models:
+                print(f"No models listed for {app_label}")
+                return
+            app_labels = [app_label]
+        else:
+            app_labels = self.django_models.keys()  # All apps
+
+        for app_label in app_labels:
+            app_name = apps.get_app_config(app_label).verbose_name
             print(f'{app_name}:')
-            for m in models:
+            for m in self.django_models[app_label]:
                 print(f'  {m}')
 
     def get_namespace(self):
-        namespace = {}
+        namespace = defaultdict(SimpleNamespace)
 
-        # Gather Django models and constants from each app
-        for app in 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)
 
-        # Additional objects to include
-        namespace['ObjectType'] = ObjectType
-        namespace['User'] = User
-
-        # Load convenience commands
-        namespace.update({
+        return {
+            **namespace,
+            'lsapps': self._lsapps,
             'lsmodels': self._lsmodels,
-        })
+        }
 
-        return namespace
+    @staticmethod
+    def get_banner_text():
+        lines = [
+            '{title} ({hostname})'.format(
+                title=bright('NetBox interactive shell'),
+                hostname=platform.node(),
+            ),
+            '{python} | {django} | {netbox}'.format(
+                python=color('green', f'Python v{platform.python_version()}'),
+                django=color('green', f'Django v{get_version()}'),
+                netbox=color('green', settings.RELEASE.name),
+            ),
+        ]
+
+        if installed_plugins := get_installed_plugins():
+            plugin_list = ', '.join([
+                color('cyan', f'{name} v{version}') for name, version in installed_plugins.items()
+            ])
+            lines.append(
+                'Plugins: {plugin_list}'.format(
+                    plugin_list=plugin_list
+                )
+            )
+
+        lines.append(
+            'lsapps() & lsmodels() will show available models. Use help(<model>) for more info.'
+        )
+
+        return '\n'.join([
+            f'### {line}' for line in lines
+        ])
 
     def handle(self, **options):
         namespace = self.get_namespace()
@@ -97,5 +157,4 @@ class Command(BaseCommand):
             readline.parse_and_bind('tab: complete')
 
         # Run interactive shell
-        shell = code.interact(banner=BANNER_TEXT, local=namespace)
-        return shell
+        return code.interact(banner=self.get_banner_text(), local=namespace)

requirements.txt

@@ -1,3 +1,4 @@
+colorama==0.4.6
 Django==5.2.4
 django-cors-headers==4.7.0
 django-debug-toolbar==5.2.0