diff --git a/src/lazy_loader/__init__.py b/src/lazy_loader/__init__.py index 23b5382..0af4e5d 100644 --- a/src/lazy_loader/__init__.py +++ b/src/lazy_loader/__init__.py @@ -21,6 +21,36 @@ threadlock = threading.Lock() +class _ShadowGuardModule(types.ModuleType): + """Module type to protect function attributes from being overwritten. + + When a function has the same name as the submodule it resides in + (e.g. a ``max_tree`` function defined in ``max_tree.py``), + importing that submodule causes the import machinery to call + ``setattr(pkg, "max_tree", )``. That updates the + package ``__dict__``, preventing ``__getattr__`` from ever + resolving the name to the function again. + + This subclass suppresses those dictionary updates (only in the + shadowing case). + + We track the set of protected names in the ``__lazy_shadowed__`` + attr. + + """ + + def __setattr__(self, name, value): + shadowed = self.__dict__.get("__lazy_shadowed__") + if ( + shadowed is not None + and name in shadowed + # Is it trying to set this attribute to the system module? + and value is sys.modules.get(f"{self.__name__}.{name}") + ): + return + super().__setattr__(name, value) + + def attach(package_name, submodules=None, submod_attrs=None): """Attach lazily loaded submodules, functions, or other attributes. @@ -92,6 +122,26 @@ def __getattr__(name): def __dir__(): return __all__.copy() + # When a function attribute has the same name as the submodule it + # resides in (e.g. `max_tree` from `max_tree.py`), importing that + # submodule makes the import machinery overwrite the parent + # package attribute with the module object, shadowing the function + # (see _ShadowGuardModule). + # + # Record affected cases and, only in those cases, swap in the + # guarding module type. + shadowed = {attr for attr, mod in attr_to_modules.items() if attr == mod} + if shadowed: + pkg = sys.modules.get(package_name) + # Only touch plain package modules (or our own wrapper) --- we + # don't want to mess with custom module classes. + if type(pkg) in (types.ModuleType, _ShadowGuardModule): + pkg.__dict__["__lazy_shadowed__"] = ( + pkg.__dict__.get("__lazy_shadowed__", set()) | shadowed + ) + if type(pkg) is types.ModuleType: + pkg.__class__ = _ShadowGuardModule + eager_import = os.environ.get("EAGER_IMPORT", "") not in ("0", "") if eager_import: for attr in set(attr_to_modules.keys()) | submodules: diff --git a/tests/test_lazy_loader.py b/tests/test_lazy_loader.py index d68537f..074395e 100644 --- a/tests/test_lazy_loader.py +++ b/tests/test_lazy_loader.py @@ -178,6 +178,17 @@ def test_attach_same_module_and_attr_name(clean_fake_pkg, eager_import): assert isinstance(some_func, types.FunctionType) +def test_attach_submodule_does_not_shadow_function(clean_fake_pkg): + # Where `some_func` is defined in module `some_func`: When + # submodule is imported before the function has been resolved, the + # import machinery tries to set the package `__dict__` to point to + # the module. We need to prevent this, otherwise we cannot + # access the function. + import tests.fake_pkg.some_func # noqa: F401 + from tests import fake_pkg + assert isinstance(fake_pkg.some_func, types.FunctionType) + + FAKE_STUB = """ from . import rank from ._gaussian import gaussian