Lazy imports in context manager.

A different (easier/safer) approach to lazy importing using a meta
import hook.

Author: vnmabus

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/nested_pkg/nested_mod_eager.py

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

lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_lazy.py

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