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.
2025-10-24 06:06:09 +00:00 · 2025-01-09 17:00:29 +01:00
parent 8c118c0e00
commit fc28550fe4
10 changed files with 374 additions and 19 deletions
--- a/tests/shell/tests.py
+++ b/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",
+        )