From 132736d885414f9e75bf22a86d7469892e6d6db5 Mon Sep 17 00:00:00 2001 From: Talley Lambert Date: Thu, 7 Jul 2022 10:04:07 -0400 Subject: [PATCH] feat: add attach_stub --- README.md | 37 ++++++++++++++++++++++ lazy_loader/__init__.py | 61 ++++++++++++++++++++++++++++++++++++- tests/fake_pkg/__init__.pyi | 1 + tests/test_lazy_loader.py | 36 ++++++++++++++++++++++ 4 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 tests/fake_pkg/__init__.pyi diff --git a/README.md b/README.md index d23854f..142309e 100644 --- a/README.md +++ b/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 diff --git a/lazy_loader/__init__.py b/lazy_loader/__init__.py index 08896b7..02f1ffe 100644 --- a/lazy_loader/__init__.py +++ b/lazy_loader/__init__.py @@ -4,6 +4,7 @@ Makes it easy to load subpackages and functions on demand. """ +import ast import importlib import importlib.util import inspect @@ -11,7 +12,7 @@ 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) diff --git a/tests/fake_pkg/__init__.pyi b/tests/fake_pkg/__init__.pyi new file mode 100644 index 0000000..d3349bb --- /dev/null +++ b/tests/fake_pkg/__init__.pyi @@ -0,0 +1 @@ +from .some_func import some_func diff --git a/tests/test_lazy_loader.py b/tests/test_lazy_loader.py index d1ee06c..b836078 100644 --- a/tests/test_lazy_loader.py +++ b/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")