diff --git a/src/lazy_loader/__init__.py b/src/lazy_loader/__init__.py index 23b5382..8448ee6 100644 --- a/src/lazy_loader/__init__.py +++ b/src/lazy_loader/__init__.py @@ -21,7 +21,12 @@ threadlock = threading.Lock() -def attach(package_name, submodules=None, submod_attrs=None): +def attach( + package_name, + submodules=None, + submod_attrs=None, + lazy_submodules=False, +): """Attach lazily loaded submodules, functions, or other attributes. Typically, modules import submodules and attributes as follows:: @@ -50,6 +55,10 @@ def attach(package_name, submodules=None, submod_attrs=None): submod_attrs : dict Dictionary of submodule -> list of attributes / functions. These attributes are imported as they are used. + lazy_submodules : bool + Whether to lazily load submodules. If set to `True`, submodules are + returned as lazy proxies. Note that attribute access from + submod_attrs will trigger the import of the submodule. Returns ------- @@ -72,7 +81,11 @@ def attach(package_name, submodules=None, submod_attrs=None): def __getattr__(name): if name in submodules: - return importlib.import_module(f"{package_name}.{name}") + submod_path = f"{package_name}.{name}" + if lazy_submodules: + return load(submod_path, suppress_warning=True) + else: + return importlib.import_module(submod_path) elif name in attr_to_modules: submod_path = f"{package_name}.{attr_to_modules[name]}" submod = importlib.import_module(submod_path) diff --git a/tests/test_lazy_loader.py b/tests/test_lazy_loader.py index d68537f..419b1b0 100644 --- a/tests/test_lazy_loader.py +++ b/tests/test_lazy_loader.py @@ -86,7 +86,8 @@ def test_lazy_import_nonbuiltins(): sp.pi -def test_lazy_attach(): +@pytest.mark.parametrize("lazy_submodules", [False, True]) +def test_lazy_attach(lazy_submodules): name = "mymod" submods = ["mysubmodule", "anothersubmodule"] myall = {"not_real_submod": ["some_var_or_func"]} @@ -96,8 +97,12 @@ def test_lazy_attach(): "name": name, "submods": submods, "myall": myall, + "lazy_submods": lazy_submodules, } - s = "__getattr__, __lazy_dir__, __all__ = attach(name, submods, myall)" + s = ( + "__getattr__, __lazy_dir__, __all__ = " + "attach(name, submods, myall, lazy_submodules=lazy_submods)" + ) exec(s, {}, locls) expected = { @@ -105,6 +110,7 @@ def test_lazy_attach(): "name": name, "submods": submods, "myall": myall, + "lazy_submods": lazy_submodules, "__getattr__": None, "__lazy_dir__": None, "__all__": None, @@ -135,9 +141,13 @@ def test_lazy_attach_noattrs(): assert all_ == sorted(submods) -def test_lazy_attach_returns_copies(): +@pytest.mark.parametrize("lazy_submodules", [False, True]) +def test_lazy_attach_returns_copies(lazy_submodules): _get, _dir, _all = lazy.attach( - __name__, ["my_submodule", "another_submodule"], {"foo": ["some_attr"]} + __name__, + ["my_submodule", "another_submodule"], + {"foo": ["some_attr"]}, + lazy_submodules=lazy_submodules, ) assert _dir() is not _dir() assert _dir() == _all