Editable installs improvements

mesonpy/_editable.py

@@ -10,6 +10,7 @@
 import importlib.abc
 import importlib.machinery
 import importlib.util
+import inspect
 import json
 import os
 import pathlib
@@ -171,19 +172,18 @@
         return MesonpyReader(name, self._tree)
 
 
-LOADERS = [
-    (ExtensionFileLoader, tuple(importlib.machinery.EXTENSION_SUFFIXES)),
-    (SourceFileLoader, tuple(importlib.machinery.SOURCE_SUFFIXES)),
-    (SourcelessFileLoader, tuple(importlib.machinery.BYTECODE_SUFFIXES)),
-]
+LOADERS = \
+    [(ExtensionFileLoader, s) for s in importlib.machinery.EXTENSION_SUFFIXES] + \
+    [(SourceFileLoader, s) for s in importlib.machinery.SOURCE_SUFFIXES] + \
+    [(SourcelessFileLoader, s) for s in importlib.machinery.BYTECODE_SUFFIXES]
 
 
 def build_module_spec(cls: type, name: str, path: str, tree: Optional[Node]) -> importlib.machinery.ModuleSpec:
     loader = cls(name, path, tree)
     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
 
 
@@ -222,13 +222,21 @@
         return dict.get(node, key)
 
 
-def walk(root: str, path: str = '') -> Iterator[pathlib.Path]:
-    with os.scandir(os.path.join(root, path)) as entries:
-        for entry in entries:
-            if entry.is_dir():
-                yield from walk(root, os.path.join(path, entry.name))
-            else:
-                yield pathlib.Path(path, entry.name)
+def walk(src: str, exclude_files: Set[str], exclude_dirs: Set[str]) -> Iterator[str]:
+    for root, dirnames, filenames in os.walk(src):
+        for name in dirnames.copy():
+            dirsrc = os.path.join(root, name)
+            relpath = os.path.relpath(dirsrc, src)
+            if relpath in exclude_dirs:
+                dirnames.remove(name)
+            # sort to process directories determninistically
+            dirnames.sort()
+        for name in sorted(filenames):
+            filesrc = os.path.join(root, name)
+            relpath = os.path.relpath(filesrc, src)
+            if relpath in exclude_files:
+                continue
+            yield relpath
 
 
 def collect(install_plan: Dict[str, Dict[str, Any]]) -> Node:
@@ -238,12 +246,43 @@
             path = pathlib.Path(target['destination'])
             if path.parts[0] in {'{py_platlib}', '{py_purelib}'}:
                 if key == 'install_subdirs' and os.path.isdir(src):
-                    for entry in walk(src):
-                        tree[(*path.parts[1:], *entry.parts)] = os.path.join(src, *entry.parts)
+                    exclude_files = {os.path.normpath(x) for x in target.get('exclude_files', [])}
+                    exclude_dirs = {os.path.normpath(x) for x in target.get('exclude_dirs', [])}
+                    for entry in walk(src, exclude_files, exclude_dirs):
+                        tree[(*path.parts[1:], *entry.split(os.sep))] = os.path.join(src, entry)
                 else:
                     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):
@@ -252,8 +291,6 @@
         self._build_cmd = cmd
         self._verbose = verbose
         self._loaders: List[Tuple[type, str]] = []
-        for loader, suffixes in LOADERS:
-            self._loaders.extend((loader, suffix) for suffix in suffixes)
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}({self._build_path!r})'
@@ -264,39 +301,15 @@
             path: Optional[Sequence[Union[bytes, str]]] = None,
             target: Optional[ModuleType] = None
     ) -> Optional[importlib.machinery.ModuleSpec]:
-        if fullname.split('.', maxsplit=1)[0] in self._top_level_modules:
-            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 self._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 self._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
+        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
+        tree = self._rebuild()
+        return find_spec(fullname, tree)
 
     @functools.lru_cache(maxsize=1)
-    def rebuild(self) -> Node:
+    def _rebuild(self) -> Node:
         # skip editable wheel lookup during rebuild: during the build
         # the module we are rebuilding might be imported causing a
         # rebuild loop.
@@ -316,6 +329,44 @@
             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/packages/imports-itself-during-build/meson.build

@@ -7,10 +7,11 @@ project('imports-itself-during-build', 'c', version: '1.0.0')
 py = import('python').find_installation()
 
 py.install_sources('pure.py')
+
 py.extension_module(
     'plat',
     'plat.c',
     install: true,
 )
 
-run_command(py, '-c', 'import pure')
+run_command(py, '-c', 'import pure', check: false)

tests/test_editable.py

@@ -4,6 +4,7 @@
 
 import os
 import pathlib
+import pkgutil
 import sys
 
 import pytest
@@ -16,12 +17,14 @@
 
 
 def test_walk(package_complex):
-    entries = set(_editable.walk(os.fspath(package_complex / 'complex')))
-    assert entries == {
+    entries = _editable.walk(
+        os.fspath(package_complex / 'complex'),
+        [os.path.normpath('more/meson.build'), os.path.normpath('more/baz.pyx')],
+        [os.path.normpath('namespace')],
+    )
+    assert {pathlib.Path(x) for x in entries} == {
         pathlib.Path('__init__.py'),
         pathlib.Path('more/__init__.py'),
-        pathlib.Path('namespace/bar.py'),
-        pathlib.Path('namespace/foo.py')
     }
 
 
@@ -189,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]