From 689fcd3a9f02697e3163a4b2d625531787513e38 Mon Sep 17 00:00:00 2001
From: Eugene Toder <eltoder@gmail.com>
Date: Fri, 22 Mar 2024 10:44:24 -0400
Subject: [PATCH] Support packages that span multiple directories

---
 src/slotscheck/discovery.py                   | 24 +++++++++----------
 .../other/implicitly_namespaced/bar.py        |  2 ++
 tests/src/conftest.py                         |  4 ++--
 tests/src/test_cli.py                         |  8 +++++++
 tests/src/test_discovery.py                   |  1 +
 5 files changed, 24 insertions(+), 15 deletions(-)
 create mode 100644 tests/examples/other/implicitly_namespaced/bar.py

diff --git a/src/slotscheck/discovery.py b/src/slotscheck/discovery.py
index 1009328..f97fe04 100644
--- a/src/slotscheck/discovery.py
+++ b/src/slotscheck/discovery.py
@@ -19,6 +19,7 @@
     Iterator,
     NamedTuple,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -171,10 +172,7 @@ def module_tree(
     if spec.submodule_search_locations is None:
         tree = Module(name)
     else:
-        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)
+        tree = _package(name, spec.submodule_search_locations)
 
     if expected_location and location != expected_location:
         raise UnexpectedImportLocation(module, expected_location, location)
@@ -188,27 +186,27 @@ def _add_namespace(tree: ModuleTree, name: ModuleNamePart) -> ModuleTree:
 
 def _submodule(m: pkgutil.ModuleInfo) -> ModuleTree:
     if m.ispkg:
-        [subdir] = m.module_finder.find_spec(
-            m.name  # type: ignore
-        ).submodule_search_locations
-        return _package(m.name, Path(subdir))
+        spec = m.module_finder.find_spec(m.name)  # type: ignore
+        assert spec is not None
+        assert spec.submodule_search_locations is not None
+        return _package(m.name, spec.submodule_search_locations)
     else:
         return Module(m.name)
 
 
-def _is_submodule(m: pkgutil.ModuleInfo, path: AbsPath) -> bool:
-    return getattr(m.module_finder, "path", "").startswith(str(path))
+def _is_submodule(m: pkgutil.ModuleInfo, paths: Tuple[str, ...]) -> bool:
+    return getattr(m.module_finder, "path", "").startswith(paths)
 
 
-def _package(name: ModuleNamePart, path: AbsPath) -> Package:
+def _package(name: ModuleNamePart, paths: Collection[str]) -> Package:
     return Package(
         name,
         frozenset(
             map(
                 _submodule,
                 filter(
-                    partial(_is_submodule, path=path),
-                    pkgutil.walk_packages([str(path)]),
+                    partial(_is_submodule, paths=tuple(paths)),
+                    pkgutil.iter_modules(paths),
                 ),
             )
         ),
diff --git a/tests/examples/other/implicitly_namespaced/bar.py b/tests/examples/other/implicitly_namespaced/bar.py
new file mode 100644
index 0000000..5d5ae96
--- /dev/null
+++ b/tests/examples/other/implicitly_namespaced/bar.py
@@ -0,0 +1,2 @@
+class C:
+    __slots__ = ()
diff --git a/tests/src/conftest.py b/tests/src/conftest.py
index 732fd4d..32ebd9f 100644
--- a/tests/src/conftest.py
+++ b/tests/src/conftest.py
@@ -13,9 +13,9 @@
 @pytest.fixture(scope="session", autouse=True)
 def add_pypath() -> Iterator[None]:
     "Add example modules to the python path"
-    sys.path.insert(0, str(EXAMPLES_DIR))
+    sys.path[:0] = [str(EXAMPLES_DIR), str(EXAMPLES_DIR / "other")]
     yield
-    sys.path.remove(str(EXAMPLES_DIR))
+    del sys.path[:2]
 
 
 @pytest.fixture(autouse=True)
diff --git a/tests/src/test_cli.py b/tests/src/test_cli.py
index 5497da4..3763baf 100644
--- a/tests/src/test_cli.py
+++ b/tests/src/test_cli.py
@@ -147,6 +147,14 @@ def test_namespaced(runner: CliRunner):
     assert result.output == "All OK!\nScanned 4 module(s), 1 class(es).\n"
 
 
+def test_implicitly_namespaced(runner: CliRunner):
+    result = runner.invoke(
+        cli, ["-m", "implicitly_namespaced"], catch_exceptions=False
+    )
+    assert result.exit_code == 0
+    assert result.output == "All OK!\nScanned 8 module(s), 2 class(es).\n"
+
+
 def test_multiple_modules(runner: CliRunner):
     result = runner.invoke(
         cli,
diff --git a/tests/src/test_discovery.py b/tests/src/test_discovery.py
index 769b795..c6502bf 100644
--- a/tests/src/test_discovery.py
+++ b/tests/src/test_discovery.py
@@ -230,6 +230,7 @@ def test_implicitly_namespaced_submodule(self):
     def test_implicitly_namespaced(self):
         assert module_tree("implicitly_namespaced", None) == make_pkg(
             "implicitly_namespaced",
+            Module("bar"),
             Module("version"),
             make_pkg("module", Module("foo"), Module("bla")),
             make_pkg("another", Module("foo")),