feat: add attach_stub function to load imports from type stubs

README.md

@@ -70,6 +70,43 @@ from .edges import (sobel, scharr, prewitt, roberts,
 
 Except that all subpackages (such as `rank`) and functions (such as `sobel`) are loaded upon access.
 
+### Lazily load subpackages and functions from type stubs
+
+Because static type checkers and IDEs will likely be unable to find your
+dynamically declared imports, you can use a [type
+stub](https://mypy.readthedocs.io/en/stable/stubs.html) (`.pyi` file) to declare
+the imports. However, if used with the above pattern, this results in code
+duplication, as you now need to declare your submodules and attributes in two places.
+
+You can infer the `submodules` and `submod_attrs` arguments (explicitly provided
+above to `lazy.attach`) from a stub  adjacent to the `.py` file by using the
+`lazy.attach_stub` function.
+
+Carrying on with the example above:
+
+The `skimage/filters/__init__.py` module would be declared as such:
+
+```python
+from ..util import lazy
+
+__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__)
+```
+
+... and the adjacent `skimage/filters/__init__.pyi` stub would contain:
+
+```python
+from . import rank
+from ._gaussian import gaussian, difference_of_gaussians
+from .edges import (sobel, scharr, prewitt, roberts,
+                    laplace, farid)
+```
+
+Note that in order for this to work, you must be sure to include the `.pyi`
+files in your package distribution. For example, with setuptools, you would need
+to [set the `package_data`
+option](https://setuptools.pypa.io/en/latest/userguide/datafiles.html#package-data)
+to include `*.pyi` files.
+
 ### Early failure
 
 With lazy loading, missing imports no longer fail upon loading the

lazy_loader/__init__.py

@@ -4,14 +4,15 @@
 
 Makes it easy to load subpackages and functions on demand.
 """
+import ast
 import importlib
 import importlib.util
 import inspect
 import os
 import sys
 import types
 
-__all__ = ["attach", "load"]
+__all__ = ["attach", "load", "attach_stub"]
 
 
 def attach(package_name, submodules=None, submod_attrs=None):
@@ -189,3 +190,61 @@ def myfunc():
     loader.exec_module(module)
 
     return module
+
+
+class _StubVisitor(ast.NodeVisitor):
+    """AST visitor to parse a stub file for submodules and submod_attrs."""
+
+    def __init__(self):
+        self._submodules = set()
+        self._submod_attrs = {}
+
+    def visit_ImportFrom(self, node: ast.ImportFrom):
+        if node.level != 1:
+            raise ValueError(
+                "Only within-module imports are supported (`from .* import`)"
+            )
+        if node.module:
+            attrs: list = self._submod_attrs.setdefault(node.module, [])
+            attrs.extend(alias.name for alias in node.names)
+        else:
+            self._submodules.update(alias.name for alias in node.names)
+
+
+def attach_stub(package_name: str, filename: str):
+    """Attach lazily loaded submodules, functions from a type stub.
+
+    This is a variant on ``attach`` that will parse a `.pyi` stub file to
+    infer ``submodules`` and ``submod_attrs``. This allows static type checkers
+    to find imports, while still providing lazy loading at runtime.
+
+    Parameters
+    ----------
+    package_name : str
+        Typically use ``__name__``.
+    filename : str
+        Path to `.py` file which has an adjacent `.pyi` file.
+        Typically use ``__file__``.
+
+    Returns
+    -------
+    __getattr__, __dir__, __all__
+        The same output as ``attach``.
+
+    Raises
+    ------
+    ValueError
+        If a stub file is not found for `filename`, or if the stubfile is formmated
+        incorrectly (e.g. if it contains an relative import from outside of the module)
+    """
+    stubfile = filename if filename.endswith("i") else f"{filename}i"
+
+    if not os.path.exists(stubfile):
+        raise ValueError(f"Cannot load imports from non-existent stub {stubfile!r}")
+
+    with open(stubfile) as f:
+        stub_node = ast.parse(f.read())
+
+    visitor = _StubVisitor()
+    visitor.visit(stub_node)
+    return attach(package_name, visitor._submodules, visitor._submod_attrs)

tests/fake_pkg/__init__.pyi

@@ -0,0 +1 @@
+from .some_func import some_func

tests/test_lazy_loader.py

@@ -108,3 +108,39 @@ def test_attach_same_module_and_attr_name():
     from fake_pkg.some_func import some_func
 
     assert isinstance(some_func, types.FunctionType)
+
+
+FAKE_STUB = """
+from . import rank
+from ._gaussian import gaussian
+from .edges import sobel, scharr, prewitt, roberts
+"""
+
+
+def test_stub_loading(tmp_path):
+    stub = tmp_path / "stub.pyi"
+    stub.write_text(FAKE_STUB)
+    _get, _dir, _all = lazy.attach_stub("my_module", str(stub))
+    expect = {"gaussian", "sobel", "scharr", "prewitt", "roberts", "rank"}
+    assert set(_dir()) == set(_all) == expect
+
+
+def test_stub_loading_parity():
+    import fake_pkg
+
+    from_stub = lazy.attach_stub(fake_pkg.__name__, fake_pkg.__file__)
+    stub_getter, stub_dir, stub_all = from_stub
+    assert stub_all == fake_pkg.__all__
+    assert stub_dir() == fake_pkg.__lazy_dir__()
+    assert stub_getter("some_func") == fake_pkg.some_func
+
+
+def test_stub_loading_errors(tmp_path):
+    stub = tmp_path / "stub.pyi"
+    stub.write_text("from ..mod import func\n")
+
+    with pytest.raises(ValueError, match="Only within-module imports are supported"):
+        lazy.attach_stub("name", str(stub))
+
+    with pytest.raises(ValueError, match="Cannot load imports from non-existent stub"):
+        lazy.attach_stub("name", "not a file")