diff --git a/lazy_loader/__init__.py b/lazy_loader/__init__.py index a306653..749333c 100644 --- a/lazy_loader/__init__.py +++ b/lazy_loader/__init__.py @@ -5,6 +5,7 @@ Makes it easy to load subpackages and functions on demand. """ import ast +import builtins import importlib import importlib.util import inspect @@ -12,6 +13,8 @@ import sys import types import warnings +from collections import defaultdict +from types import SimpleNamespace __all__ = ["attach", "load", "attach_stub"] @@ -257,11 +260,14 @@ 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()) @@ -269,3 +275,76 @@ def attach_stub(package_name: str, filename: str): 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 diff --git a/lazy_loader/tests/fake_pkg_magic/__init__.py b/lazy_loader/tests/fake_pkg_magic/__init__.py new file mode 100644 index 0000000..9e20af8 --- /dev/null +++ b/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 diff --git a/lazy_loader/tests/fake_pkg_magic/nested_pkg/__init__.py b/lazy_loader/tests/fake_pkg_magic/nested_pkg/__init__.py new file mode 100644 index 0000000..439dd46 --- /dev/null +++ b/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 diff --git a/lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_eager.py b/lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_eager.py new file mode 100644 index 0000000..e69de29 diff --git a/lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_lazy.py b/lazy_loader/tests/fake_pkg_magic/nested_pkg/nested_mod_lazy.py new file mode 100644 index 0000000..e69de29 diff --git a/lazy_loader/tests/fake_pkg_magic/some_func.py b/lazy_loader/tests/fake_pkg_magic/some_func.py new file mode 100644 index 0000000..10e99ed --- /dev/null +++ b/lazy_loader/tests/fake_pkg_magic/some_func.py @@ -0,0 +1,3 @@ +def some_func(): + """Function with same name as submodule.""" + pass diff --git a/lazy_loader/tests/fake_pkg_magic/some_mod.py b/lazy_loader/tests/fake_pkg_magic/some_mod.py new file mode 100644 index 0000000..6f432cd --- /dev/null +++ b/lazy_loader/tests/fake_pkg_magic/some_mod.py @@ -0,0 +1,2 @@ +class SomeClass: + pass \ No newline at end of file