Add context manager functionality using hooks

lazy_loader/__init__.py

@@ -5,16 +5,30 @@
 Makes it easy to load subpackages and functions on demand.
 """
 import ast
+import builtins
 import importlib
 import importlib.util
 import inspect
 import os
 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,15 +271,72 @@ 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())
 
     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)

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

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

lazy_loader/tests/fake_pkg_magic/some_func.py

@@ -0,0 +1,3 @@
+def some_func():
+    """Function with same name as submodule."""
+    pass

lazy_loader/tests/fake_pkg_magic/some_mod.py

@@ -0,0 +1,2 @@
+class SomeClass:
+    pass

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)