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.

2 files changed, 199 insertions, 8 deletions

diff --git a/tests/shell/models.py b/tests/shell/models.py
new file mode 100644
index 0000000000..85b40bf205
--- /dev/null
+++ b/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)
diff --git a/tests/shell/tests.py b/tests/shell/tests.py
index ca823f6290..49528cca8e 100644
--- 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",
+        )