Add context manager functionality

lazy_loader/__init__.py

@@ -5,13 +5,16 @@
 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 collections import defaultdict
+from types import SimpleNamespace
 
 __all__ = ["attach", "load", "attach_stub"]
 
@@ -257,15 +260,91 @@ 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)
+
+
+PLACEHOLDER = object()
+
+
+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 __init__(self):
+        self.imports = []
+        self.submodules = []
+        self.submod_attrs = defaultdict(list)
+
+    def __enter__(self):
+        # Prevent normal importing
+        self.import_fun = builtins.__import__
+        builtins.__import__ = self._my_import
+        return self
+
+    def __exit__(self, type, value, tb):
+        # Restore importing
+        builtins.__import__ = self.import_fun
+
+        last_frame = inspect.currentframe().f_back
+
+        # Remove imported things
+        for submod in self.submodules:
+            mod = submod.partition(".")[0]
+            del last_frame.f_globals[mod]
+
+        for mod, attr_list in self.submod_attrs.items():
+            for attr in attr_list:
+                del last_frame.f_globals[attr]
+
+        # Inject the outputs into the module globals
+        package_name = last_frame.f_globals["__name__"]
+        g, d, a = attach(
+            package_name,
+            self.submodules,
+            self.submod_attrs,
+        )
+
+        last_frame.f_globals["__getattr__"] = g
+        last_frame.f_globals["__dir__"] = d
+        last_frame.f_globals["__all__"] = a
+
+    def _my_import(self, name, globals=None, locals=None, fromlist=(), level=0):
+        builtins.__import__ = self.import_fun
+        self.imports.append(
+            {"name": name, "fromlist": fromlist, "level": level}
+        )
+        if fromlist is None:
+            raise NotImplementedError(
+                "Absolute imports are not currently "
+                "supported by lazy_loader."
+            )
+        elif name == "":
+            self.submodules.extend(fromlist)
+        else:
+            self.submod_attrs[name].extend(fromlist)
+
+        builtins.__import__ = self._my_import
+
+        if fromlist:
+            return SimpleNamespace(**{k: PLACEHOLDER for k in fromlist})
+        return PLACEHOLDER

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