Check __spec__.has_location + refactor

Author: loic-simon

Lib/_pyrepl/_module_completer.py

@@ -16,6 +16,7 @@
 TYPE_CHECKING = False
 
 if TYPE_CHECKING:
+    from types import ModuleType
     from typing import Any, Iterable, Iterator, Mapping
 
 
@@ -110,17 +111,7 @@ def _find_modules(self, path: str, prefix: str) -> list[str]:
         modules: Iterable[pkgutil.ModuleInfo]
         imported_module = sys.modules.get(path.split('.')[0])
         if imported_module:
-            # Module already imported: only look in its location,
-            # even if a module with the same name would be higher in path
-            imported_path = (imported_module.__spec__
-                             and imported_module.__spec__.origin)
-            if not imported_path:
-                # Module imported but no spec/origin: propose no suggestions
-                return []
-            if os.path.basename(imported_path) == "__init__.py":  # package
-                imported_path = os.path.dirname(imported_path)
-            import_location = os.path.dirname(imported_path)
-            modules = list(pkgutil.iter_modules([import_location]))
+            modules = self._find_already_imported_module_specs(imported_module)
         else:
             modules = self.global_cache
 
@@ -145,6 +136,32 @@ def _is_stdlib_module(self, module_info: pkgutil.ModuleInfo) -> bool:
         return (isinstance(module_info.module_finder, FileFinder)
                 and module_info.module_finder.path == self._stdlib_path)
 
+    def _find_already_imported_module_specs(self, imported_module: ModuleType) -> list[pkgutil.ModuleInfo]:
+        # Module already imported: only look in its location,
+        # even if a module with the same name would be higher in path
+        module_location = self._get_module_location(imported_module)
+        if not module_location:
+            # If we cannot find the module source, propose no suggestions
+            return []
+        import_location = os.path.dirname(module_location)
+        return list(pkgutil.iter_modules([import_location]))
+
+    def _get_module_location(self, imported_module: ModuleType) -> str | None:
+        spec = imported_module.__spec__
+        if not spec:
+            return None
+        if not spec.has_location:
+            if spec.origin == "frozen":  # See Tools/build/freeze_modules.py
+                return os.path.join(self._stdlib_path, f"{spec.name}.py")
+            return None
+        if not spec.origin:
+            return None
+        if imported_module.__package__:
+            # Package: the module location is the parent folder
+            return os.path.dirname(spec.origin)
+        else:
+            return spec.origin
+
     def is_suggestion_match(self, module_name: str, prefix: str) -> bool:
         if prefix:
             return module_name.startswith(prefix)

Lib/test/test_pyrepl/test_pyrepl.py

@@ -972,6 +972,7 @@ def test_import_completions(self):
             ("from importlib import mac\t\n", "from importlib import machinery"),
             ("from importlib import res\t\n", "from importlib import resources"),
             ("from importlib.res\t import a\t\n", "from importlib.resources import abc"),
+            ("from __phello__ import s\t\n", "from __phello__ import spam"),  # frozen module
         )
         for code, expected in cases:
             with self.subTest(code=code):
@@ -1132,6 +1133,14 @@ def test_already_imported_stdlib_module_no_other_suggestions(self):
             output = reader.readline()
             self.assertEqual(output, "import collections.abc")
 
+    @patch.dict(sys.modules)
+    def test_already_imported_frozen_module(self):
+        importlib.import_module("__phello__")
+        events = code_to_events("from __phello__ import s\t\n")
+        reader = self.prepare_reader(events, namespace={})
+        output = reader.readline()
+        self.assertEqual(output, "from __phello__ import spam")
+
     def test_already_imported_custom_module_no_other_suggestions(self):
         with (tempfile.TemporaryDirectory() as _dir1,
               tempfile.TemporaryDirectory() as _dir2,
@@ -1187,14 +1196,16 @@ def test_already_imported_module_without_origin_or_spec(self):
         with (tempfile.TemporaryDirectory() as _dir1,
               patch.object(sys, "path", [_dir1, *sys.path])):
             dir1 = pathlib.Path(_dir1)
-            for mod in ("no_origin", "no_spec"):
+            for mod in ("no_origin", "not_has_location", "no_spec"):
                 (dir1 / mod).mkdir()
                 (dir1 / mod / "__init__.py").touch()
                 (dir1 / mod / "foo.py").touch()
                 module = importlib.import_module(mod)
                 assert module.__spec__
                 if mod == "no_origin":
                     module.__spec__.origin = None
+                elif mod == "not_has_location":
+                    module.__spec__.has_location = False
                 else:
                     module.__spec__ = None
                 events = code_to_events(f"import {mod}.\t\n")