Guard against import machinery masking out attributes

src/lazy_loader/__init__.py

@@ -21,6 +21,37 @@
 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", <submodule>)``.  That updates the
+    package ``__dict__``, preventing ``__getattr__`` from ever
+    resolving the name to the function again. The same problem occurs
+    when ``x`` is defined in ``x/sub.py``.
+
+    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 +123,29 @@ def __getattr__(name):
     def __dir__():
         return __all__.copy()
 
+    # When a function has the same name as a module the import
+    # machinery needs to load along the way to accessing it
+    # (e.g. `max_tree` from `max_tree.py`, or `x` from `x/sub.py`), a
+    # side-effect of it loading that module is overwriting the package
+    # attribute (so it points to the module, i.e. to `max_tree` or `x`
+    # the module), 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.split(".")[0]
+    }
+    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:

tests/fake_pkg_submodule/__init__.py

@@ -0,0 +1,3 @@
+import lazy_loader as lazy
+
+__getattr__, __dir__, __all__ = lazy.attach(__name__, submod_attrs={"x.sub": ["x"]})

tests/fake_pkg_submodule/x/sub.py

@@ -0,0 +1,2 @@
+def x():
+    """Function with the same name as the package it lives under."""

tests/test_lazy_loader.py

@@ -13,9 +13,9 @@
 @pytest.fixture
 def clean_fake_pkg():
     yield
-    sys.modules.pop("tests.fake_pkg.some_func", None)
-    sys.modules.pop("tests.fake_pkg", None)
-    sys.modules.pop("tests", None)
+    for name in [*sys.modules]:
+        if name == "tests" or name.startswith("tests.fake_pkg"):
+            del sys.modules[name]
 
 
 @pytest.mark.parametrize("attempt", [1, 2])
@@ -178,6 +178,30 @@ 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)
+
+
+def test_attach_subpackage_does_not_shadow_function(clean_fake_pkg):
+    # Where function `x` is defined in `x/sub.py`, i.e. a module
+    # nested inside a subpackage that shares its name with the
+    # function: importing `x.sub` causes the import machinery to set
+    # the package attribute `x` to the `x` subpackage, which would
+    # otherwise shadow the `x` function on subsequent access.
+    from tests import fake_pkg_submodule
+
+    assert isinstance(fake_pkg_submodule.x, types.FunctionType)
+    assert isinstance(fake_pkg_submodule.x, types.FunctionType)
+
+
 FAKE_STUB = """
 from . import rank
 from ._gaussian import gaussian