diff --git a/lazy_loader/__init__.py b/lazy_loader/__init__.py index a306653..72a0786 100644 --- a/lazy_loader/__init__.py +++ b/lazy_loader/__init__.py @@ -5,6 +5,7 @@ Makes it easy to load subpackages and functions on demand. """ import ast +import builtins import importlib import importlib.util import inspect @@ -12,9 +13,22 @@ import sys import types import warnings +from contextvars import ContextVar +from importlib.abc import MetaPathFinder +from importlib.machinery import ModuleSpec +from typing import Sequence __all__ = ["attach", "load", "attach_stub"] +inside_context_manager: ContextVar[bool] = ContextVar( + 'inside_context_manager', + default=False, +) +searching_spec: ContextVar[bool] = ContextVar( + 'searching_spec', + default=False, +) + def attach(package_name, submodules=None, submod_attrs=None): """Attach lazily loaded submodules, functions, or other attributes. @@ -257,11 +271,14 @@ def attach_stub(package_name: str, filename: str): incorrectly (e.g. if it contains an relative import from outside of the module) """ stubfile = ( - filename if filename.endswith("i") else f"{os.path.splitext(filename)[0]}.pyi" + filename if filename.endswith("i") + else f"{os.path.splitext(filename)[0]}.pyi" ) if not os.path.exists(stubfile): - raise ValueError(f"Cannot load imports from non-existent stub {stubfile!r}") + raise ValueError( + f"Cannot load imports from non-existent stub {stubfile!r}", + ) with open(stubfile) as f: stub_node = ast.parse(f.read()) @@ -269,3 +286,57 @@ def attach_stub(package_name: str, filename: str): visitor = _StubVisitor() visitor.visit(stub_node) return attach(package_name, visitor._submodules, visitor._submod_attrs) + + +class LazyFinder(MetaPathFinder): + + def find_spec( + self, + fullname: str, + path: Sequence[str] | None, + target: types.ModuleType | None = ..., + / , + ) -> ModuleSpec | None: + if not inside_context_manager.get(): + # We are not in context manager, delegate to normal import + return None + + if searching_spec.get(): + # We are searching for the loader, so we should continue the search + return None + + searching_spec.set(True) + spec = importlib.util.find_spec(fullname) + searching_spec.set(False) + + if spec is None: + raise ModuleNotFoundError(f"No module named '{fullname}'") + + spec.loader = importlib.util.LazyLoader(spec.loader) + + return spec + + +sys.meta_path.insert(0, LazyFinder()) + + +class lazy_imports: + """ + Context manager that will block imports and make them lazy. + + >>> import lazy_loader + >>> with lazy_loader.lazy_imports(): + >>> from ._mod import some_func + + """ + + def __enter__(self): + # Prevent normal importing + if inside_context_manager.get(): + raise ValueError("Nested lazy_imports not allowed.") + inside_context_manager.set(True) + return self + + def __exit__(self, type, value, tb): + # Restore normal importing + inside_context_manager.set(False) diff --git a/lazy_loader/tests/fake_pkg_magic/__init__.py b/lazy_loader/tests/fake_pkg_magic/__init__.py new file mode 100644 index 0000000..9e20af8 --- /dev/null +++ b/lazy_loader/tests/fake_pkg_magic/__init__.py @@ -0,0 +1,5 @@ +import lazy_loader as lazy + +with lazy.lazy_imports(): + from .some_func import some_func + from . import some_mod, nested_pkg diff --git a/lazy_loader/tests/fake_pkg_magic/nested_pkg/__init__.py b/lazy_loader/tests/fake_pkg_magic/nested_pkg/__init__.py new file mode 100644 index 0000000..439dd46 --- /dev/null +++ b/lazy_loader/tests/fake_pkg_magic/nested_pkg/__init__.py @@ -0,0 +1,6 @@ +import lazy_loader as lazy + +from . import nested_mod_eager + +with lazy.lazy_imports(): + from . import nested_mod_lazy diff --git a/lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_eager.py b/lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_eager.py new file mode 100644 index 0000000..e69de29 diff --git a/lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_lazy.py b/lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_lazy.py new file mode 100644 index 0000000..e69de29 diff --git a/lazy_loader/tests/fake_pkg_magic/some_func.py b/lazy_loader/tests/fake_pkg_magic/some_func.py new file mode 100644 index 0000000..10e99ed --- /dev/null +++ b/lazy_loader/tests/fake_pkg_magic/some_func.py @@ -0,0 +1,3 @@ +def some_func(): + """Function with same name as submodule.""" + pass diff --git a/lazy_loader/tests/fake_pkg_magic/some_mod.py b/lazy_loader/tests/fake_pkg_magic/some_mod.py new file mode 100644 index 0000000..6f432cd --- /dev/null +++ b/lazy_loader/tests/fake_pkg_magic/some_mod.py @@ -0,0 +1,2 @@ +class SomeClass: + pass \ No newline at end of file diff --git a/lazy_loader/tests/test_lazy_loader_magic.py b/lazy_loader/tests/test_lazy_loader_magic.py new file mode 100644 index 0000000..266d903 --- /dev/null +++ b/lazy_loader/tests/test_lazy_loader_magic.py @@ -0,0 +1,75 @@ +import importlib +import sys +import types + +import lazy_loader as lazy +import pytest + + +def test_lazy_import_basics(): + with lazy.lazy_imports(): + import math + + with pytest.raises(ImportError): + with lazy.lazy_imports(): + import anything_not_real + + # Now test that accessing attributes does what it should + assert math.sin(math.pi) == pytest.approx(0, 1e-6) + + +def test_lazy_import_subpackages(): + with lazy.lazy_imports(): + import html.parser as hp + assert "html" in sys.modules + assert type(sys.modules["html"]) == type(pytest) + assert isinstance(hp, importlib.util._LazyModule) + assert "html.parser" in sys.modules + assert sys.modules["html.parser"] == hp + + +def test_lazy_import_impact_on_sys_modules(): + with lazy.lazy_imports(): + import math + + with pytest.raises(ImportError): + with lazy.lazy_imports(): + import anything_not_real + + assert isinstance(math, types.ModuleType) + assert "math" in sys.modules + assert "anything_not_real" not in sys.modules + + # only do this if numpy is installed + pytest.importorskip("numpy") + with lazy.lazy_imports(): + import numpy as np + assert isinstance(np, types.ModuleType) + assert "numpy" in sys.modules + + np.pi # trigger load of numpy + + assert isinstance(np, types.ModuleType) + assert "numpy" in sys.modules + + +def test_lazy_import_nonbuiltins(): + with lazy.lazy_imports(): + import numpy as np + import scipy as sp + + assert np.sin(np.pi) == pytest.approx(0, 1e-6) + + +def test_attach_same_module_and_attr_name(): + from lazy_loader.tests import fake_pkg_magic + + # Grab attribute twice, to ensure that importing it does not + # override function by module + assert isinstance(fake_pkg_magic.some_func, types.FunctionType) + assert isinstance(fake_pkg_magic.some_func, types.FunctionType) + + # Ensure imports from submodule still work + from lazy_loader.tests.fake_pkg_magic.some_func import some_func + + assert isinstance(some_func, types.FunctionType)