Fix importing of namespace packages and built-in modules

pyproject.toml

@@ -29,14 +29,14 @@ tomli = {version = ">=0.2.6,<3.0.0", python = "<3.11"}
 [tool.poetry.dev-dependencies]
 flake8 = "^7.0"
 isort = "^5.13.2"
-mypy = "^1.9"
+mypy = {version = "^1.9", extras = ["mypyc"]}
 pytest = "^8.1.1"
 black = "^24"
 pytest-cov = "^4.1.0"
 pytest-mock = "^3.12.0"
 typing_extensions = ">=4.1,<5"
-# Only actually needed as an example library whose __init__ is an extension.
-pydantic = "<3"
+# Used as an example of an extension module.
+ujson = "^5.9.0"
 
 [tool.poetry.scripts]
 slotscheck = "slotscheck.cli:root"

src/slotscheck/discovery.py

@@ -4,11 +4,6 @@
 import pkgutil
 from dataclasses import dataclass, field, replace
 from functools import partial, reduce
-from importlib._bootstrap_external import (  # type: ignore[import-not-found]
-    _NamespaceLoader,
-)
-from importlib.abc import FileLoader
-from importlib.machinery import ExtensionFileLoader
 from importlib.util import find_spec
 from inspect import isclass
 from itertools import chain, takewhile
@@ -149,7 +144,7 @@ def merge(self, other: ModuleTree) -> ModuleTree:
 class UnexpectedImportLocation(Exception):
     module: ModuleName
     expected: AbsPath
-    actual: AbsPath
+    actual: Optional[AbsPath]
 
 
 @add_slots
@@ -170,26 +165,16 @@ def module_tree(
         return FailedImport(module, e)
     if spec is None:
         raise ModuleNotFoundError(f"No module named '{module}'", name=module)
-    loader = spec.loader
     *namespaces, name = module.split(".")
-    location: AbsPath
+    location = Path(spec.origin) if spec.has_location and spec.origin else None
     tree: ModuleTree
-    if isinstance(loader, (FileLoader, ExtensionFileLoader)):
-        assert isinstance(loader.path, str)  # type: ignore[union-attr]
-        location = Path(loader.path)  # type: ignore[union-attr]
-        tree = (
-            _package(name, location.parent)
-            if loader.is_package(module)
-            else Module(name)
-        )
-    elif isinstance(loader, _NamespaceLoader):
-        assert len(loader._path._path) == 1
-        location = Path(loader._path._path[0])
-        tree = _package(name, location)
-    elif module == "builtins":
-        return Module(module)
+    if spec.submodule_search_locations is None:
+        tree = Module(name)
     else:
-        raise NotImplementedError(f"Unsupported module loader type: {loader}")
+        assert len(spec.submodule_search_locations) == 1
+        pkg_location = Path(spec.submodule_search_locations[0])
+        location = location or pkg_location
+        tree = _package(name, pkg_location)
 
     if expected_location and location != expected_location:
         raise UnexpectedImportLocation(module, expected_location, location)

tests/__init__.py

@@ -0,0 +1 @@
+# This file is needed for the overrides in mypy.ini to work.

tests/examples/gc.py

@@ -0,0 +1 @@
+# This module is shadowed by the builtin module.

tests/src/test_cli.py

@@ -103,9 +103,9 @@ def test_builtins(runner: CliRunner):
 
 
 def test_extension(runner: CliRunner):
-    result = runner.invoke(cli, ["-m", "_pickle"])
+    result = runner.invoke(cli, ["-m", "ujson"])
     assert result.exit_code == 0
-    assert result.output == ("All OK!\nScanned 1 module(s), 5 class(es).\n")
+    assert result.output == ("All OK!\nScanned 1 module(s), 1 class(es).\n")
 
 
 def test_success_verbose(runner: CliRunner):

tests/src/test_discovery.py

@@ -221,27 +221,26 @@ def test_namespaced(self):
             "namespaced", make_pkg("module", Module("foo"), Module("bla"))
         )
 
-    def test_implicitly_namspaced_submodule(self):
+    def test_implicitly_namespaced_submodule(self):
         assert module_tree("implicitly_namespaced.module", None) == make_pkg(
             "implicitly_namespaced",
             make_pkg("module", Module("foo"), Module("bla")),
         )
 
-    def test_namespace_loader(self):
-        import implicitly_namespaced.module  # type: ignore  # noqa
-
+    def test_implicitly_namespaced(self):
         assert module_tree("implicitly_namespaced", None) == make_pkg(
             "implicitly_namespaced",
             Module("version"),
             make_pkg("module", Module("foo"), Module("bla")),
             make_pkg("another", Module("foo")),
         )
 
-    def test_builtin(self):
-        assert module_tree("builtins", None) == Module("builtins")
+    @pytest.mark.parametrize("module", ["_ast", "builtins", "gc", "sys"])
+    def test_builtin(self, module: str):
+        assert module_tree(module, None) == Module(module)
 
     def test_extension(self):
-        assert module_tree("_elementtree", None) == Module("_elementtree")
+        assert module_tree("ujson", None) == Module("ujson")
 
     def test_import_causes_base_exception_no_strict_imports(self, mocker):
         assert module_tree(
@@ -255,9 +254,9 @@ def test_import_error(self, mocker):
         ) == FailedImport("broken.submodule", mocker.ANY)
 
     def test_extension_package(self):
-        tree = module_tree("pydantic", None)
+        tree = module_tree("mypyc", None)
         assert isinstance(tree, Package)
-        assert len(tree.content) > 20
+        assert len(tree.content) > 10
 
     def test_module(self):
         assert module_tree(
@@ -282,6 +281,14 @@ def test_unexpected_location(self):
             EXAMPLES_DIR / "module_misc/a/b/c.py",
         )
 
+    def test_shadowed_by_builtin(self):
+        with pytest.raises(UnexpectedImportLocation) as exc:
+            module_tree("gc", expected_location=EXAMPLES_DIR / "gc.py")
+
+        assert exc.value == UnexpectedImportLocation(
+            "gc", EXAMPLES_DIR / "gc.py", None
+        )
+
     def test_pyc_file(self):
         assert module_tree("compiled", None) == make_pkg(
             "compiled", Module("foo"), Module("bar")