Auto-import stdlib modules

Author: loic-simon

Lib/_pyrepl/_module_completer.py

@@ -157,6 +157,16 @@ def _find_attributes(self, path: str, prefix: str) -> tuple[list[str], Completio
         if not imported_module:
             if path in self._failed_imports:  # Do not propose to import again
                 return [], None
+            root = path.split(".")[0]
+            mod_info = next((m for m in self.global_cache if m.name == root),
+                             None)
+            if mod_info and self._is_stdlib_module(mod_info):
+                # Stdlib module: auto-import (no risk of dangerous side-effect)
+                try:
+                    imported_module = importlib.import_module(path)
+                except Exception:
+                    sys.modules.pop(path, None)  # Clean half-imported module
+        if not imported_module:
             return [], self._get_import_completion_action(path)
         try:
             module_attributes = dir(imported_module)

Lib/test/test_pyrepl/test_pyrepl.py

@@ -1140,6 +1140,7 @@ def test_attribute_completion_module_on_demand(self):
             (dir / "pack" / "foo.py").touch()
             (dir / "pack" / "bar.py").touch()
             (dir / "pack" / "baz.py").touch()
+            sys.modules.pop("graphlib", None)  # test modules may have been imported by previous tests
             with patch.object(sys, "path", [_dir, *sys.path]):
                 cases = (
                     # needs 2 tabs to import (show prompt, then import)
@@ -1160,6 +1161,8 @@ def test_attribute_completion_module_on_demand(self):
                     ("from pack import b\t\n", "from pack import ba", set()),
                     ("from pack import b\t\t\n", "from pack import ba", set()),
                     ("from pack import b\t\t\t\n", "from pack import ba", {"pack"}),
+                    # stdlib modules are automatically imported
+                    ("from graphlib import T\t\n", "from graphlib import TopologicalSorter", {"graphlib"}),
                 )
                 for code, expected, expected_imports in cases:
                     with self.subTest(code=code), patch.dict(sys.modules):
@@ -1370,6 +1373,8 @@ def test_suggestions_and_messages(self) -> None:
             (dir / "pack").mkdir()
             (dir / "pack" / "__init__.py").write_text("foo = 1; bar = 2;")
             (dir / "pack" / "bar.py").touch()
+            sys.modules.pop("graphlib", None)  # test modules may have been imported by previous tests
+            sys.modules.pop("compression.zstd", None)
             with patch.object(sys, "path", [_dir, *sys.path]):
                 cases = (
                     # no match != not an import
@@ -1391,6 +1396,9 @@ def test_suggestions_and_messages(self) -> None:
                     ("from pack import ", (["bar", "foo"], None), set()),
                     ("from pack.bar import ", ([], (_prompt, None)), {"pack.bar"}),
                     ("from pack.bar import ", ([], None), set()),
+                    # stdlib = auto-imported
+                    ("from graphlib import T", (["TopologicalSorter"], None), {"graphlib"}),
+                    ("from compression.zstd import c", (["compress"], None), {"compression.zstd"}),
                 )
                 completer = ModuleCompleter()
                 for i, (code, expected, expected_imports) in enumerate(cases):
@@ -1410,8 +1418,6 @@ def test_suggestions_and_messages(self) -> None:
 
                         new_imports = sys.modules.keys() - _imported
                         self.assertSetEqual(new_imports, expected_imports)
-                        for mod in new_imports:
-                            self.addCleanup(sys.modules.pop, mod)
 
 class TestHardcodedSubmodules(TestCase):
     @patch.dict(sys.modules)