ENH: set __path__ for module loaded in editable installation

Set the __path__ module attribute to a placeholder path.  Because how
editable installs are implemented, this cannot correspond to the
filesystem path for the package.  And add a sys.path_hook that
recognizes this paths and returns a path loader that implements
iter_modules(). This allows pkgutil.iter_packages() to work properly.

Fixes #557.
Fixes #568.

Author: dnicolodi

mesonpy/_editable.py

@@ -10,6 +10,7 @@
 import importlib.abc
 import importlib.machinery
 import importlib.util
+import inspect
 import json
 import os
 import pathlib
@@ -182,7 +183,7 @@ def build_module_spec(cls: type, name: str, path: str, tree: Optional[Node]) ->
     spec = importlib.machinery.ModuleSpec(name, loader, origin=path)
     spec.has_location = True
     if loader.is_package(name):
-        spec.submodule_search_locations = []
+        spec.submodule_search_locations = [os.path.join(__file__, name)]
     return spec
 
 
@@ -253,6 +254,35 @@ def collect(install_plan: Dict[str, Dict[str, Any]]) -> Node:
                     tree[path.parts[1:]] = src
     return tree
 
+def find_spec(fullname: str, tree: Node) -> Optional[importlib.machinery.ModuleSpec]:
+    namespace = False
+    parts = fullname.split('.')
+
+    # look for a package
+    package = tree.get(tuple(parts))
+    if isinstance(package, Node):
+        for loader, suffix in LOADERS:
+            src = package.get('__init__' + suffix)
+            if isinstance(src, str):
+                return build_module_spec(loader, fullname, src, package)
+        else:
+            namespace = True
+
+    # look for a module
+    for loader, suffix in LOADERS:
+        src = tree.get((*parts[:-1], parts[-1] + suffix))
+        if isinstance(src, str):
+            return build_module_spec(loader, fullname, src, None)
+
+    # namespace
+    if namespace:
+        spec = importlib.machinery.ModuleSpec(fullname, None, is_package=True)
+        assert isinstance(spec.submodule_search_locations, list)  # make mypy happy
+        spec.submodule_search_locations.append(os.path.join(__file__, fullname))
+        return spec
+
+    return None
+
 
 class MesonpyMetaFinder(importlib.abc.MetaPathFinder):
     def __init__(self, names: Set[str], path: str, cmd: List[str], verbose: bool = False):
@@ -271,40 +301,12 @@ def find_spec(
             path: Optional[Sequence[Union[bytes, str]]] = None,
             target: Optional[ModuleType] = None
     ) -> Optional[importlib.machinery.ModuleSpec]:
-
-        if fullname.split('.', maxsplit=1)[0] not in self._top_level_modules:
+        if fullname.split('.', 1)[0] not in self._top_level_modules:
             return None
-
         if self._build_path in os.environ.get(MARKER, '').split(os.pathsep):
             return None
-
-        namespace = False
         tree = self._rebuild()
-        parts = fullname.split('.')
-
-        # look for a package
-        package = tree.get(tuple(parts))
-        if isinstance(package, Node):
-            for loader, suffix in LOADERS:
-                src = package.get('__init__' + suffix)
-                if isinstance(src, str):
-                    return build_module_spec(loader, fullname, src, package)
-            else:
-                namespace = True
-
-        # look for a module
-        for loader, suffix in LOADERS:
-            src = tree.get((*parts[:-1], parts[-1] + suffix))
-            if isinstance(src, str):
-                return build_module_spec(loader, fullname, src, None)
-
-        # namespace
-        if namespace:
-            spec = importlib.machinery.ModuleSpec(fullname, None)
-            spec.submodule_search_locations = []
-            return spec
-
-        return None
+        return find_spec(fullname, tree)
 
     @functools.lru_cache(maxsize=1)
     def _rebuild(self) -> Node:
@@ -327,6 +329,44 @@ def _rebuild(self) -> Node:
             install_plan = json.load(f)
         return collect(install_plan)
 
+    def _path_hook(self, path: str) -> MesonpyPathFinder:
+        if os.altsep:
+            path.replace(os.altsep, os.sep)
+        path, _, key = path.rpartition(os.sep)
+        if path == __file__:
+            tree = self._rebuild()
+            node = tree.get(tuple(key.split('.')))
+            if isinstance(node, Node):
+                return MesonpyPathFinder(node)
+        raise ImportError
+
+
+class MesonpyPathFinder(importlib.abc.PathEntryFinder):
+    def __init__(self, tree: Node):
+        self._tree = tree
+
+    def find_spec(self, fullname: str, target: Optional[ModuleType] = None) -> Optional[importlib.machinery.ModuleSpec]:
+        return find_spec(fullname, self._tree)
+
+    def iter_modules(self, prefix: str) -> Iterator[Tuple[str, bool]]:
+        yielded = set()
+        for name, node in self._tree.items():
+            modname = inspect.getmodulename(name)
+            if modname == '__init__' or modname in yielded:
+                continue
+            if isinstance(node, Node):
+                modname = name
+                for _, suffix in LOADERS:
+                    src = node.get('__init__' + suffix)
+                    if isinstance(src, str):
+                        yielded.add(modname)
+                        yield prefix + modname, True
+            elif modname and '.' not in modname:
+                yielded.add(modname)
+                yield prefix + modname, False
+
 
 def install(names: Set[str], path: str, cmd: List[str], verbose: bool) -> None:
-    sys.meta_path.insert(0, MesonpyMetaFinder(names, path, cmd, verbose))
+    finder = MesonpyMetaFinder(names, path, cmd, verbose)
+    sys.meta_path.insert(0, finder)
+    sys.path_hooks.insert(0, finder._path_hook)

tests/packages/complex/complex/more/baz.pyx

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2023 The meson-python developers
+#
+# SPDX-License-Identifier: MIT
+
+def answer():
+    return 42

tests/packages/complex/complex/more/meson.build

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: 2024 The meson-python developers
+#
+# SPDX-License-Identifier: MIT
+
+py.extension_module(
+    'baz',
+    'baz.pyx',
+    install: true,
+    subdir: 'complex/more',
+)

tests/packages/complex/foo.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2024 The meson-python developers
+#
+# SPDX-License-Identifier: MIT
+
+def foo():
+    return True

tests/packages/complex/meson.build

@@ -15,6 +15,36 @@ endif
 
 py = import('python').find_installation()
 
-install_subdir('complex', install_dir: py.get_install_dir(pure: false))
+py.install_sources(
+    'move.py',
+    subdir: 'complex/more',
+    pure: false,
+)
 
-py.extension_module('test', 'test.pyx', install: true, subdir: 'complex')
+install_data(
+    'foo.py',
+    rename: 'bar.py',
+    install_dir: py.get_install_dir(pure: false) / 'complex',
+)
+
+install_subdir(
+    'complex',
+    install_dir: py.get_install_dir(pure: false),
+    exclude_files: ['more/meson.build', 'more/baz.pyx'],
+)
+
+py.extension_module(
+    'test',
+    'test.pyx',
+    install: true,
+    subdir: 'complex',
+)
+
+py.extension_module(
+    'baz',
+    'complex/more/baz.pyx',
+    install: true,
+    subdir: 'complex/more',
+)
+
+subdir('complex/more')

tests/packages/complex/move.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2024 The meson-python developers
+#
+# SPDX-License-Identifier: MIT
+
+def test():
+    return True

tests/test_editable.py

@@ -4,6 +4,7 @@
 
 import os
 import pathlib
+import pkgutil
 import sys
 
 import pytest
@@ -191,3 +192,37 @@ def test_editble_reentrant(venv, editable_imports_itself_during_build):
         assert venv.python('-c', 'import plat; print(plat.data())').strip() == 'DEF'
     finally:
         path.write_text(code)
+
+
+def test_editable_pkgutils_walk_packages(package_complex, tmp_path):
+    # build a package in a temporary directory
+    mesonpy.Project(package_complex, tmp_path)
+
+    finder = _editable.MesonpyMetaFinder({'complex'}, os.fspath(tmp_path), ['ninja'])
+
+    try:
+        # install editable hooks
+        sys.meta_path.insert(0, finder)
+        sys.path_hooks.insert(0, finder._path_hook)
+
+        import complex
+        packages = {m.name for m in pkgutil.walk_packages(complex.__path__, complex.__name__ + '.')}
+        assert packages == {
+            'complex.bar',
+            'complex.more',
+            'complex.more.baz',
+            'complex.more.move',
+            'complex.test',
+        }
+
+        from complex import namespace
+        packages = {m.name for m in pkgutil.walk_packages(namespace.__path__, namespace.__name__ + '.')}
+        assert packages == {
+            'complex.namespace.bar',
+            'complex.namespace.foo',
+        }
+
+    finally:
+        # remove hooks
+        del sys.meta_path[0]
+        del sys.path_hooks[0]