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

@@ -2,7 +2,9 @@ import os
 import select
 import sys
 import traceback
+from collections import defaultdict
 
+from django.apps import apps
 from django.core.management import BaseCommand, CommandError
 from django.utils.datastructures import OrderedSet
 
@@ -26,6 +28,11 @@ class Command(BaseCommand):
                 "variable and ~/.pythonrc.py script."
             ),
         )
+        parser.add_argument(
+            "--no-imports",
+            action="store_true",
+            help="Disable automatic imports of models.",
+        )
         parser.add_argument(
             "-i",
             "--interface",
@@ -47,18 +54,27 @@ 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"], options["no_imports"]
+            ),
+        )
 
     def bpython(self, options):
         import bpython
 
-        bpython.embed()
+        bpython.embed(
+            self.get_and_report_namespace(options["verbosity"], options["no_imports"])
+        )
 
     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"], options["no_imports"]
+        )
 
         # We want to honor both $PYTHONSTARTUP and .pythonrc.py, so follow system
         # conventions and get $PYTHONSTARTUP first then .pythonrc.py.
@@ -111,10 +127,68 @@ class Command(BaseCommand):
         # Start the interactive interpreter.
         code.interact(local=imported_objects)
 
+    def get_and_report_namespace(self, verbosity, no_imports=False):
+        if no_imports:
+            return {}
+
+        namespace = self.get_namespace()
+
+        if verbosity < 1:
+            return namespace
+
+        amount = len(namespace)
+        msg = f"{amount} objects imported automatically"
+
+        if verbosity < 2:
+            self.stdout.write(f"{msg} (use -v 2 for details).", self.style.SUCCESS)
+            return namespace
+
+        imports_by_module = defaultdict(list)
+        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__")
+            ):
+                imports_by_module[obj.__module__].append(obj_name)
+            if not hasattr(obj, "__module__") and hasattr(obj, "__name__"):
+                tokens = obj.__name__.split(".")
+                if obj_name in tokens:
+                    module = ".".join(t for t in tokens if t != obj_name)
+                    imports_by_module[module].append(obj_name)
+
+        import_string = "\n".join(
+            [
+                f"  from {module} import {objects}"
+                for module, imported_objects in imports_by_module.items()
+                if (objects := ", ".join(imported_objects))
+            ]
+        )
+
+        try:
+            import isort
+        except ImportError:
+            pass
+        else:
+            import_string = isort.code(import_string)
+
+        self.stdout.write(
+            f"{msg}, including:\n\n{import_string}", self.style.SUCCESS, ending="\n\n"
+        )
+
+        return namespace
+
+    def get_namespace(self):
+        apps_models = apps.get_models()
+        namespace = {}
+        for model in reversed(apps_models):
+            if model.__module__:
+                namespace[model.__name__] = model
+        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 +198,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,57 @@
+======================================
+How to customize the ``shell`` command
+======================================
+
+The Django :djadmin:`shell` is an interactive Python environment that provides
+access to models and settings, making it useful for testing code, experimenting
+with queries, and interacting with application data.
+
+Customizing the :djadmin:`shell` command allows adding extra functionality or
+pre-loading specific modules. To do this, create a new management command that subclasses
+``django.core.management.commands.shell.Command`` and overrides the existing
+``shell`` management command. For more details, refer to the guide on
+:ref:`overriding commands <overriding-commands>`.
+
+.. _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

@@ -58,8 +58,9 @@ Other guides
 
    auth-remote-user
    csrf
-   custom-management-commands
    custom-file-storage
+   custom-management-commands
+   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.
+By default, the :djadmin:`shell` command automatically imports the models from
+your :setting:`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,19 @@ Mails the email addresses specified in :setting:`ADMINS` using
 
 Starts the Python interactive interpreter.
 
+All models from installed apps are automatically imported into the shell
+environment. Models from apps listed earlier in :setting:`INSTALLED_APPS` take
+precedence. For a ``--verbosity`` of 2 or higher, the automatically imported
+objects will be listed. To disable automatic importing entirely, use the
+``--no-imports`` flag.
+
+See the guide on :ref:`customizing this behaviour
+<customizing-shell-auto-imports>` to add or remove automatic imports.
+
+.. versionchanged:: 5.2
+
+    Automatic models import 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,26 @@ 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 models from
+all installed apps. You can view further details of the imported objects by
+setting the ``--verbosity`` flag to 2 or more:
+
+.. code-block:: pycon
+
+    $ python -Wall manage.py shell --verbosity=2
+    6 objects imported automatically, including:
+
+      from django.contrib.admin.models import LogEntry
+      from django.contrib.auth.models import Group, Permission, User
+      from django.contrib.contenttypes.models import ContentType
+      from django.contrib.sessions.models import Session
+
+This :ref:`behavior can be customized <customizing-shell-auto-imports>` to add
+or remove automatic 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,25 @@ import unittest
 from unittest import mock
 
 from django import __version__
+from django.contrib.auth.models import Group, Permission, User
+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 .models import Marker, Phone
 
 
 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 +87,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, "no_imports": False})
 
-        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 +108,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, "no_imports": False})
 
-        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 +128,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, "no_imports": False})
 
-        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 +141,166 @@ 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.
+
+
+class ShellCommandAutoImportsTestCase(SimpleTestCase):
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_get_namespace(self):
+        namespace = shell.Command().get_namespace()
+
+        self.assertEqual(
+            namespace,
+            {
+                "Marker": Marker,
+                "Phone": Phone,
+                "ContentType": ContentType,
+                "Group": Group,
+                "Permission": Permission,
+                "User": User,
+            },
+        )
+
+    @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"
+
+        winner_article = Article
+
+        class Article(models.Model):
+            class Meta:
+                app_label = "shell"
+
+        with mock.patch("django.apps.apps.get_models", return_value=apps.get_models()):
+            namespace = shell.Command().get_namespace()
+            self.assertEqual(namespace, {"Article": winner_article})
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_get_namespace_overridden(self):
+        class TestCommand(shell.Command):
+            def get_namespace(self):
+                from django.urls.base import resolve, reverse
+
+                return {
+                    **super().get_namespace(),
+                    "resolve": resolve,
+                    "reverse": reverse,
+                }
+
+        namespace = TestCommand().get_namespace()
+
+        self.assertEqual(
+            namespace,
+            {
+                "resolve": resolve,
+                "reverse": reverse,
+                "Marker": Marker,
+                "Phone": Phone,
+                "ContentType": ContentType,
+                "Group": Group,
+                "Permission": Permission,
+                "User": User,
+            },
+        )
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_no_imports_flag(self):
+        for verbosity in (0, 1, 2, 3):
+            with self.subTest(verbosity=verbosity), captured_stdout() as stdout:
+                namespace = shell.Command().get_and_report_namespace(
+                    verbosity=verbosity, no_imports=True
+                )
+            self.assertEqual(namespace, {})
+            self.assertEqual(stdout.getvalue().strip(), "")
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_verbosity_zero(self):
+        with captured_stdout() as stdout:
+            cmd = shell.Command()
+            namespace = cmd.get_and_report_namespace(verbosity=0)
+        self.assertEqual(namespace, cmd.get_namespace())
+        self.assertEqual(stdout.getvalue().strip(), "")
+
+    @override_settings(
+        INSTALLED_APPS=["shell", "django.contrib.auth", "django.contrib.contenttypes"]
+    )
+    def test_verbosity_one(self):
+        with captured_stdout() as stdout:
+            cmd = shell.Command()
+            namespace = cmd.get_and_report_namespace(verbosity=1)
+        self.assertEqual(namespace, cmd.get_namespace())
+        self.assertEqual(
+            stdout.getvalue().strip(),
+            "6 objects imported automatically (use -v 2 for details).",
+        )
+
+    @override_settings(INSTALLED_APPS=["shell", "django.contrib.contenttypes"])
+    @mock.patch.dict(sys.modules, {"isort": None})
+    def test_message_with_stdout_listing_objects_with_isort_not_installed(self):
+        class TestCommand(shell.Command):
+            def get_namespace(self):
+                class MyClass:
+                    pass
+
+                constant = "constant"
+
+                return {
+                    **super().get_namespace(),
+                    "MyClass": MyClass,
+                    "constant": constant,
+                }
+
+        with captured_stdout() as stdout:
+            TestCommand().get_and_report_namespace(verbosity=2)
+
+        self.assertEqual(
+            stdout.getvalue().strip(),
+            "5 objects imported automatically, including:\n\n"
+            "  from django.contrib.contenttypes.models import ContentType\n"
+            "  from shell.models import Phone, Marker",
+        )
+
+    @override_settings(INSTALLED_APPS=["shell", "django.contrib.contenttypes"])
+    def test_message_with_stdout_listing_objects_with_isort(self):
+        sorted_imports = (
+            "  from shell.models import Marker, Phone\n\n"
+            "  from django.contrib.contenttypes.models import ContentType"
+        )
+        mock_isort_code = mock.Mock(code=mock.MagicMock(return_value=sorted_imports))
+
+        class TestCommand(shell.Command):
+            def get_namespace(self):
+                class MyClass:
+                    pass
+
+                constant = "constant"
+
+                return {
+                    **super().get_namespace(),
+                    "MyClass": MyClass,
+                    "constant": constant,
+                }
+
+        with (
+            mock.patch.dict(sys.modules, {"isort": mock_isort_code}),
+            captured_stdout() as stdout,
+        ):
+            TestCommand().get_and_report_namespace(verbosity=2)
+
+        self.assertEqual(
+            stdout.getvalue().strip(),
+            "5 objects imported automatically, including:\n\n"
+            "  from shell.models import Marker, Phone\n\n"
+            "  from django.contrib.contenttypes.models import ContentType",
+        )