diff --git a/lazy_loader/__init__.py b/lazy_loader/__init__.py index 02f1ffe..6533054 100644 --- a/lazy_loader/__init__.py +++ b/lazy_loader/__init__.py @@ -4,6 +4,8 @@ Makes it easy to load subpackages and functions on demand. """ +from __future__ import annotations + import ast import importlib import importlib.util @@ -11,8 +13,10 @@ import os import sys import types +from contextlib import contextmanager +from importlib.util import LazyLoader -__all__ = ["attach", "load", "attach_stub"] +__all__ = ["attach", "load", "attach_stub", "lazy_loader"] def attach(package_name, submodules=None, submod_attrs=None): @@ -248,3 +252,35 @@ 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: + @classmethod + def find_spec(cls, name, path, target=None) -> LazyLoader | None: + """Finds a spec with every other Finder in sys.meta_path, + and, if found, wraps it in LazyLoader to defer loading. + """ + non_lazy_finders = (f for f in sys.meta_path if f is not cls) + for finder in non_lazy_finders: + spec = finder.find_spec(name, path, target) + if spec is not None: + spec.loader = LazyLoader(spec.loader) + break + return spec + + +@contextmanager +def lazy_import(): + """A context manager to defer imports until first access. + + >>> with lazy_import(): + ... import math # lazy + ... + >>> math.inf # executes the math module. + inf + + lazy_import inserts, and then removes, LazyFinder to the start of sys.meta_path. + """ + sys.meta_path.insert(0, LazyFinder) + yield + sys.meta_path.pop(0) diff --git a/tests/fake_pkg/__init__.py b/tests/fake_pkg/__init__.py index 540fd73..f271945 100644 --- a/tests/fake_pkg/__init__.py +++ b/tests/fake_pkg/__init__.py @@ -1,5 +1,4 @@ -import lazy_loader as lazy +from lazy_loader import lazy_import -__getattr__, __lazy_dir__, __all__ = lazy.attach( - __name__, submod_attrs={"some_func": ["some_func"]} -) +with lazy_import(): + from . import some_func # noqa: F401 diff --git a/tests/test_lazy_loader.py b/tests/test_lazy_loader.py index b836078..42fff7d 100644 --- a/tests/test_lazy_loader.py +++ b/tests/test_lazy_loader.py @@ -4,41 +4,51 @@ import pytest import lazy_loader as lazy +from lazy_loader import lazy_import -def test_lazy_import_basics(): - math = lazy.load("math") - anything_not_real = lazy.load("anything_not_real") +def test_import_builtin(): + with lazy_import(): + import math # Now test that accessing attributes does what it should assert math.sin(math.pi) == pytest.approx(0, 1e-6) - # poor-mans pytest.raises for testing errors on attribute access - try: - anything_not_real.pi - assert False # Should not get here - except ModuleNotFoundError: - pass - assert isinstance(anything_not_real, lazy.DelayedImportErrorModule) - # see if it changes for second access - try: - anything_not_real.pi - assert False # Should not get here - except ModuleNotFoundError: - pass - - -def test_lazy_import_impact_on_sys_modules(): - math = lazy.load("math") - anything_not_real = lazy.load("anything_not_real") + + +def test_import_error(): + with pytest.raises(ModuleNotFoundError): + with lazy_import(): + import anything_not_real # noqa: F401 + + +def test_import_nonbuiltins(): + pytest.importorskip("numpy") + + with lazy_import(): + import numpy as np + + assert np.sin(np.pi) == pytest.approx(0, 1e-6) + + +def test_builtin_is_in_sys_modules(): + with lazy_import(): + import math assert isinstance(math, types.ModuleType) assert "math" in sys.modules - assert isinstance(anything_not_real, lazy.DelayedImportErrorModule) - assert "anything_not_real" not in sys.modules + math.pi # trigger load of math + + assert isinstance(math, types.ModuleType) + assert "math" in sys.modules + + +def test_non_builtin_is_in_sys_modules(): # only do this if numpy is installed pytest.importorskip("numpy") - np = lazy.load("numpy") + with lazy_import(): + import numpy as np + assert isinstance(np, types.ModuleType) assert "numpy" in sys.modules @@ -48,25 +58,6 @@ def test_lazy_import_impact_on_sys_modules(): assert "numpy" in sys.modules -def test_lazy_import_nonbuiltins(): - sp = lazy.load("scipy") - np = lazy.load("numpy") - if isinstance(sp, lazy.DelayedImportErrorModule): - try: - sp.pi - assert False - except ModuleNotFoundError: - pass - elif isinstance(np, lazy.DelayedImportErrorModule): - try: - np.sin(np.pi) - assert False - except ModuleNotFoundError: - pass - else: - assert np.sin(sp.pi) == pytest.approx(0, 1e-6) - - def test_lazy_attach(): name = "mymod" submods = ["mysubmodule", "anothersubmodule"] @@ -101,8 +92,8 @@ def test_attach_same_module_and_attr_name(): # Grab attribute twice, to ensure that importing it does not # override function by module - assert isinstance(fake_pkg.some_func, types.FunctionType) - assert isinstance(fake_pkg.some_func, types.FunctionType) + assert isinstance(fake_pkg.some_func, types.ModuleType) + assert isinstance(fake_pkg.some_func, types.ModuleType) # Ensure imports from submodule still work from fake_pkg.some_func import some_func