Fixed #35515 -- Added automatic model imports to shell management command.

Thanks to Bhuvnesh Sharma and Adam Johnson for mentoring this Google
Summer of Code 2024 project. Thanks to Sarah Boyce, David Smith, Jacob
Walls and Natalia Bidart for reviews.

django/core/management/commands/shell.py

@@ -3,6 +3,7 @@ import select
 import sys
 import traceback
 
+from django.apps import apps
 from django.core.management import BaseCommand, CommandError
 from django.utils.datastructures import OrderedSet
 
@@ -47,18 +48,21 @@ class Command(BaseCommand):
     def ipython(self, options):
         from IPython import start_ipython
 
-        start_ipython(argv=[])
+        start_ipython(
+            argv=[],
+            user_ns=self.get_and_report_namespace(options["verbosity"]),
+        )
 
     def bpython(self, options):
         import bpython
 
-        bpython.embed()
+        bpython.embed(self.get_and_report_namespace(options["verbosity"]))
 
     def python(self, options):
         import code
 
         # Set up a dictionary to serve as the environment for the shell.
-        imported_objects = {}
+        imported_objects = self.get_and_report_namespace(options["verbosity"])
 
         # We want to honor both $PYTHONSTARTUP and .pythonrc.py, so follow system
         # conventions and get $PYTHONSTARTUP first then .pythonrc.py.
@@ -111,10 +115,66 @@ class Command(BaseCommand):
         # Start the interactive interpreter.
         code.interact(local=imported_objects)
 
+    def get_and_report_namespace(self, verbosity):
+        namespace = self.get_namespace()
+
+        if verbosity >= 1:
+            self.stdout.write(
+                f"{len(namespace)} objects imported automatically",
+                self.style.SUCCESS,
+            )
+        if verbosity >= 2:
+            imports_by_module = {}
+            imports_by_alias = {}
+            for obj_name, obj in namespace.items():
+                if hasattr(obj, "__module__") and (
+                    (hasattr(obj, "__qualname__") and obj.__qualname__.find(".") == -1)
+                    or not hasattr(obj, "__qualname__")
+                ):
+                    module = obj.__module__
+                    collected_imports = imports_by_module.get(module, [])
+                    imports_by_module[module] = collected_imports + [obj_name]
+                if not hasattr(obj, "__module__") and hasattr(obj, "__name__"):
+                    tokens = obj.__name__.split(".")
+                    if obj_name in tokens:
+                        tokens.remove(obj_name)
+                        module = ".".join(tokens)
+                        collected_imports = imports_by_module.get(module, [])
+                        imports_by_module[module] = collected_imports + [obj_name]
+                    else:
+                        module = ".".join(tokens)
+                        imports_by_alias[module] = obj_name
+
+            for module, imported_objects in imports_by_module.items():
+                self.stdout.write(
+                    f"from {module} import {', '.join(imported_objects)}",
+                    self.style.SUCCESS,
+                )
+            for module, alias in imports_by_alias.items():
+                self.stdout.write(f"import {module} as {alias}", self.style.SUCCESS)
+
+        return namespace
+
+    def get_namespace(self):
+        apps_models = apps.get_models()
+        apps_models_modules = [
+            (app.models_module, app.label)
+            for app in apps.get_app_configs()
+            if app.models_module is not None
+        ]
+        namespace = {}
+        for model in reversed(apps_models):
+            if model.__module__:
+                namespace[model.__name__] = model
+        for app_models_module, app_label in apps_models_modules:
+            if f"{app_label}_models" not in namespace:
+                namespace[f"{app_label}_models"] = app_models_module
+        return namespace
+
     def handle(self, **options):
         # Execute the command and exit.
         if options["command"]:
-            exec(options["command"], globals())
+            exec(options["command"], {**globals(), **self.get_namespace()})
             return
 
         # Execute stdin if it has anything to read and exit.
@@ -124,7 +184,7 @@ class Command(BaseCommand):
             and not sys.stdin.isatty()
             and select.select([sys.stdin], [], [], 0)[0]
         ):
-            exec(sys.stdin.read(), globals())
+            exec(sys.stdin.read(), {**globals(), **self.get_namespace()})
             return
 
         available_shells = (

docs/howto/custom-management-commands.txt

@@ -157,6 +157,8 @@ Testing
 Information on how to test custom management commands can be found in the
 :ref:`testing docs <topics-testing-management-commands>`.
 
+.. _overriding-commands:
+
 Overriding commands
 ===================
 

docs/howto/custom-shell.txt

@@ -0,0 +1,58 @@
+======================================
+How to customize the ``shell`` command
+======================================
+
+The Django :djadmin:`shell` is an interactive Python environment that provides
+access to models and settings. It is useful for testing code, experimenting
+with queries, and interacting with application data.
+
+Customizing the :djadmin:`shell` command can improve workflow by adding
+functionality or pre-loading specific modules. To do this, first, create a new
+custom management command which subclasses
+``django.core.management.commands.shell.Command``, overriding the existing
+``shell`` management command. See the guide on :ref:`overriding commands
+<overriding-commands>` for more details.
+
+.. _customizing-shell-auto-imports:
+
+Customize automatic imports
+===========================
+
+.. versionadded:: 5.2
+
+To customize the automatic import behavior of the :djadmin:`shell` management
+command, override the ``get_namespace()`` method. For example:
+
+.. code-block:: python
+    :caption: ``polls/management/commands/shell.py``
+
+    from django.core.management.commands import shell
+
+
+    class Command(shell.Command):
+        def get_namespace(self):
+            from django.urls.base import resolve, reverse
+
+            return {
+                **super().get_namespace(),
+                "resolve": resolve,
+                "reverse": reverse,
+            }
+
+The above customization adds :func:`~django.urls.resolve` and
+:func:`~django.urls.reverse` to the default namespace, which includes all
+models from all apps. These two functions will then be available when the
+shell opens, without a manual import statement.
+
+If you prefer to not have models automatically imported, create a custom
+``get_namespace()`` that excludes the ``super().get_namespace()`` call:
+
+.. code-block:: python
+    :caption: ``polls/management/commands/shell.py``
+
+    from django.core.management.commands import shell
+
+
+    class Command(shell.Command):
+        def get_namespace(self):
+            return {}

docs/howto/index.txt

@@ -60,6 +60,7 @@ Other guides
    csrf
    custom-management-commands
    custom-file-storage
+   custom-shell
 
 .. seealso::
 

docs/intro/tutorial02.txt

@@ -347,13 +347,13 @@ API Django gives you. To invoke the Python shell, use this command:
 We're using this instead of simply typing "python", because :file:`manage.py`
 sets the :envvar:`DJANGO_SETTINGS_MODULE` environment variable, which gives
 Django the Python import path to your :file:`mysite/settings.py` file.
+Additionally, the :djadmin:`shell` command automatically imports the models
+from your installed apps.
 
 Once you're in the shell, explore the :doc:`database API </topics/db/queries>`:
 
 .. code-block:: pycon
 
-    >>> from polls.models import Choice, Question  # Import the model classes we just wrote.
-
     # No questions are in the system yet.
     >>> Question.objects.all()
     <QuerySet []>
@@ -443,8 +443,6 @@ Save these changes and start a new Python interactive shell by running
 
 .. code-block:: pycon
 
-    >>> from polls.models import Choice, Question
-
     # Make sure our __str__() addition worked.
     >>> Question.objects.all()
     <QuerySet [<Question: What's up?>]>

docs/intro/tutorial05.txt

@@ -150,7 +150,6 @@ whose date lies in the future:
 
     >>> import datetime
     >>> from django.utils import timezone
-    >>> from polls.models import Question
     >>> # create a Question instance with pub_date 30 days in the future
     >>> future_question = Question(pub_date=timezone.now() + datetime.timedelta(days=30))
     >>> # was it published recently?

docs/ref/django-admin.txt

@@ -1065,6 +1065,21 @@ Mails the email addresses specified in :setting:`ADMINS` using
 
 Starts the Python interactive interpreter.
 
+All models from all installed apps are automatically imported. See the guide on
+:ref:`customizing this behaviour <customizing-shell-auto-imports>` to add or
+remove imports.
+
+Models from apps listed earlier in :setting:`INSTALLED_APPS` take precedence.
+You can access overridden models from later, lower-precedence apps through its
+models module, which is made available as ``<app_label>_models``, for example
+``biblio_models.Book`` to access a model called ``Book`` in an app called
+``biblio``. For a ``--verbosity`` of 2 and above, all imported functions and
+classes are shown.
+
+.. versionchanged:: 5.2
+
+    Automatic importing of models was added.
+
 .. django-admin-option:: --interface {ipython,bpython,python}, -i {ipython,bpython,python}
 
 Specifies the shell to use. By default, Django will use IPython_ or bpython_ if

docs/releases/5.2.txt

@@ -31,6 +31,31 @@ and only officially support the latest release of each series.
 What's new in Django 5.2
 ========================
 
+Automatic models import in the ``shell``
+----------------------------------------
+
+The :djadmin:`shell` management command now automatically imports all models
+from all installed apps.
+
+You can view all imported functions and classes setting the ``--verbosity``
+flag to 2 or more.
+
+.. code-block:: shell
+
+    $ python manage.py shell --verbosity=2
+    10 objects imported automatically
+    from django.contrib.sessions.models import Session
+    from django.contrib.contenttypes.models import ContentType
+    from django.contrib.auth.models import User, Group, Permission
+    from django.contrib.admin.models import LogEntry
+    import django.contrib.admin.models as admin_models
+    import django.contrib.auth.models as auth_models
+    import django.contrib.contenttypes.models as contenttypes_models
+    import django.contrib.sessions.models as sessions_models
+
+This :ref:`behavior can be customized <customizing-shell-auto-imports>` to add
+extra imports or remove some or all of the default imports.
+
 Composite Primary Keys
 ----------------------
 

tests/shell/models.py

@@ -0,0 +1,9 @@
+from django.db import models
+
+
+class Marker(models.Model):
+    pass
+
+
+class Phone(models.Model):
+    name = models.CharField(max_length=50)

tests/shell/tests.py

@@ -3,14 +3,27 @@ import unittest
 from unittest import mock
 
 from django import __version__
+from django.contrib.auth import models as auth_model_module
+from django.contrib.auth.models import Group, Permission, User
+from django.contrib.contenttypes import models as contenttypes_model_module
+from django.contrib.contenttypes.models import ContentType
 from django.core.management import CommandError, call_command
 from django.core.management.commands import shell
+from django.db import models
 from django.test import SimpleTestCase
-from django.test.utils import captured_stdin, captured_stdout
+from django.test.utils import (
+    captured_stdin,
+    captured_stdout,
+    isolate_apps,
+    override_settings,
+)
+from django.urls.base import resolve, reverse
+
+from . import models as shell_models
 
 
 class ShellCommandTestCase(SimpleTestCase):
-    script_globals = 'print("__name__" in globals())'
+    script_globals = 'print("__name__" in globals() and "Phone" in globals())'
     script_with_inline_function = (
         "import django\ndef f():\n    print(django.__version__)\nf()"
     )
@@ -76,9 +89,12 @@ class ShellCommandTestCase(SimpleTestCase):
         mock_ipython = mock.Mock(start_ipython=mock.MagicMock())
 
         with mock.patch.dict(sys.modules, {"IPython": mock_ipython}):
-            cmd.ipython({})
+            cmd.ipython({"verbosity": 0})
 
-        self.assertEqual(mock_ipython.start_ipython.mock_calls, [mock.call(argv=[])])
+        self.assertEqual(
+            mock_ipython.start_ipython.mock_calls,
+            [mock.call(argv=[], user_ns=cmd.get_and_report_namespace(0))],
+        )
 
     @mock.patch("django.core.management.commands.shell.select.select")  # [1]
     @mock.patch.dict("sys.modules", {"IPython": None})
@@ -94,9 +110,11 @@ class ShellCommandTestCase(SimpleTestCase):
         mock_bpython = mock.Mock(embed=mock.MagicMock())
 
         with mock.patch.dict(sys.modules, {"bpython": mock_bpython}):
-            cmd.bpython({})
+            cmd.bpython({"verbosity": 0})
 
-        self.assertEqual(mock_bpython.embed.mock_calls, [mock.call()])
+        self.assertEqual(
+            mock_bpython.embed.mock_calls, [mock.call(cmd.get_and_report_namespace(0))]
+        )
 
     @mock.patch("django.core.management.commands.shell.select.select")  # [1]
     @mock.patch.dict("sys.modules", {"bpython": None})
@@ -112,9 +130,12 @@ class ShellCommandTestCase(SimpleTestCase):
         mock_code = mock.Mock(interact=mock.MagicMock())
 
         with mock.patch.dict(sys.modules, {"code": mock_code}):
-            cmd.python({"no_startup": True})
+            cmd.python({"verbosity": 0, "no_startup": True})
 
-        self.assertEqual(mock_code.interact.mock_calls, [mock.call(local={})])
+        self.assertEqual(
+            mock_code.interact.mock_calls,
+            [mock.call(local=cmd.get_and_report_namespace(0))],
+        )
 
     # [1] Patch select to prevent tests failing when the test suite is run
     # in parallel mode. The tests are run in a subprocess and the subprocess's
@@ -122,3 +143,156 @@ class ShellCommandTestCase(SimpleTestCase):
     # returns EOF and so select always shows that sys.stdin is ready to read.
     # This causes problems because of the call to select.select() toward the
     # end of shell's handle() method.
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_get_namespace(self):
+        cmd = shell.Command()
+        namespace = cmd.get_namespace()
+
+        self.assertEqual(
+            namespace,
+            {
+                "Marker": shell_models.Marker,
+                "Phone": shell_models.Phone,
+                "ContentType": ContentType,
+                "Group": Group,
+                "Permission": Permission,
+                "User": User,
+                "auth_models": auth_model_module,
+                "contenttypes_models": contenttypes_model_module,
+                "shell_models": shell_models,
+            },
+        )
+
+    @override_settings(INSTALLED_APPS=["basic", "shell"])
+    @isolate_apps("basic", "shell", kwarg_name="apps")
+    def test_get_namespace_precedence(self, apps):
+        class Article(models.Model):
+            class Meta:
+                app_label = "basic"
+
+        article_basic = Article
+
+        class Article(models.Model):
+            class Meta:
+                app_label = "shell"
+
+        article_shell = Article
+
+        cmd = shell.Command()
+
+        with mock.patch("django.apps.apps.get_models", return_value=apps.get_models()):
+            namespace = cmd.get_namespace()
+            self.assertIn(article_basic, namespace.values())
+            self.assertNotIn(article_shell, namespace.values())
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_overridden_get_namespace(self):
+        class Command(shell.Command):
+            def get_namespace(self):
+                from django.urls.base import resolve, reverse
+
+                return {
+                    **super().get_namespace(),
+                    "resolve": resolve,
+                    "reverse": reverse,
+                }
+
+        cmd = Command()
+        namespace = cmd.get_namespace()
+
+        self.assertEqual(
+            namespace,
+            {
+                "resolve": resolve,
+                "reverse": reverse,
+                "Marker": shell_models.Marker,
+                "Phone": shell_models.Phone,
+                "ContentType": ContentType,
+                "Group": Group,
+                "Permission": Permission,
+                "User": User,
+                "auth_models": auth_model_module,
+                "contenttypes_models": contenttypes_model_module,
+                "shell_models": shell_models,
+            },
+        )
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_message_with_stdout(self):
+        with captured_stdout() as stdout:
+            cmd = shell.Command()
+            cmd.get_and_report_namespace(verbosity=1)
+        self.assertEqual(stdout.getvalue().strip(), "9 objects imported automatically")
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_message_with_no_stdout(self):
+        with captured_stdout() as stdout:
+            cmd = shell.Command()
+            cmd.get_and_report_namespace(verbosity=0)
+        self.assertEqual(stdout.getvalue().strip(), "")
+
+    @override_settings(INSTALLED_APPS=["shell", "django.contrib.contenttypes"])
+    def test_message_with_stdout_listing_objects(self):
+        class Command(shell.Command):
+            def get_namespace(self):
+                class MyClass:
+                    pass
+
+                constant = "constant"
+
+                return {
+                    **super().get_namespace(),
+                    "MyClass": MyClass,
+                    "constant": constant,
+                }
+
+        with captured_stdout() as stdout:
+            cmd = Command()
+            cmd.get_and_report_namespace(verbosity=2)
+
+        self.assertEqual(
+            stdout.getvalue().strip(),
+            "7 objects imported automatically\n"
+            "from django.contrib.contenttypes.models import ContentType\n"
+            "from shell.models import Phone, Marker\n"
+            "import shell.models as shell_models\n"
+            "import django.contrib.contenttypes.models as contenttypes_models",
+        )
+
+    @override_settings(
+        INSTALLED_APPS=[
+            "shell",
+            "django.contrib.auth",
+            "django.contrib.messages",
+            "django.contrib.staticfiles",
+        ]
+    )
+    def test_shell_has_no_empty_models_objects(self):
+        cmd = shell.Command()
+        namespace = cmd.get_namespace()
+
+        self.assertNotIn("messages_models", namespace)
+        self.assertNotIn("staticfiles_models", namespace)
+        self.assertIsNotNone(namespace["auth_models"])
+        self.assertIsNotNone(namespace["shell_models"])
+
+    @override_settings(INSTALLED_APPS=["shell"])
+    @isolate_apps("shell", kwarg_name="apps")
+    def test_app_has_model_with_name_equal_module_name(self, apps):
+        class shell_models(models.Model):
+            pass
+
+        cmd = shell.Command()
+
+        with mock.patch("django.apps.apps.get_models", return_value=apps.get_models()):
+            namespace = cmd.get_namespace()
+            self.assertEqual(namespace["shell_models"], shell_models)