lazy_import as a context manager

lazy_loader/__init__.py

@@ -4,15 +4,19 @@
 
 Makes it easy to load subpackages and functions on demand.
 """
+from __future__ import annotations
+
 import ast
 import importlib
 import importlib.util
 import inspect
 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)

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

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