feat: support for namespace packages

src/lazy_imports_lite/_loader.py

@@ -38,7 +38,7 @@ def scan_distributions():
         metadata = dist.metadata
 
         if metadata is None:
-            continue
+            continue  # pragma: no cover
 
         if metadata["Keywords"] is None:
             continue
@@ -56,14 +56,21 @@ def _top_level_declared(dist):
 def _top_level_inferred(dist):
     files = dist.files
     if files is None:
-        return {}
+        return {}  # pragma: no cover
 
-    return {
-        f.parts[0] if len(f.parts) > 1 else f.with_suffix("").name
+    parts = {
+        f.parts[:-1] if len(f.parts) > 1 else f.with_suffix("").name
         for f in files
         if f.suffix == ".py"
     }
 
+    is_namespace = min(len(p) for p in parts) == 2
+
+    if is_namespace:
+        return {".".join(p) for p in parts if len(p) == 2}
+    else:
+        return {".".join(p) for p in parts if len(p) == 1}
+
 
 class LazyLoader(importlib.abc.Loader, importlib.machinery.PathFinder):
     def find_spec(self, fullname, path=None, target=None):
@@ -83,8 +90,11 @@ def find_spec(self, fullname, path=None, target=None):
             return None  # pragma: no cover
 
         name = spec.name.split(".")[0]
+        namespace_name = ".".join(spec.name.split(".")[:2])
 
-        if name in enabled_packages and spec.origin.endswith(".py"):
+        if (
+            name in enabled_packages or namespace_name in enabled_packages
+        ) and spec.origin.endswith(".py"):
             spec.loader = self
             return spec
 

tests/test_loader.py

@@ -22,7 +22,7 @@ def write_files(dir, content):
 
 
 @contextmanager
-def package(name, content, extra_config=""):
+def package(name, content, extra_config="", lazy_imports_enabled=True):
     content = {
         "pyproject.toml": f"""
 
@@ -32,7 +32,7 @@ def package(name, content, extra_config=""):
 
 [project]
 name="{name}"
-keywords=["lazy-imports-lite-enabled"]
+keywords=[{'"lazy-imports-lite-enabled"' if lazy_imports_enabled else ""}]
 version="0.0.1"
 """
         + extra_config,
@@ -559,3 +559,67 @@ def foo():
         ),
         normal_stderr=snapshot(""),
     )
+
+
+def test_namespace_package():
+    check_script(
+        [
+            package(
+                "test-pck-a",
+                {
+                    "source/test_pck/a/__init__.py": """\
+
+def foo():
+    print("foo")
+
+""",
+                },
+                extra_config="""
+[tool.hatch.build.targets.wheel]
+packages = ["source/test_pck"]
+""",
+            ),
+            package(
+                "test-pck-b",
+                {
+                    "source/test_pck/b/__init__.py": """\
+
+def foo():
+    print("foo")
+
+""",
+                },
+                extra_config="""
+[tool.hatch.build.targets.wheel]
+packages = ["source/test_pck"]
+""",
+                lazy_imports_enabled=False,
+            ),
+        ],
+        """\
+import test_pck
+import test_pck.a
+import test_pck.b
+
+print(type(test_pck.__spec__.loader))
+print(type(test_pck.a.__spec__.loader))
+print(type(test_pck.b.__spec__.loader))
+
+""",
+        transformed_stdout=snapshot(
+            """\
+<class '_frozen_importlib_external.NamespaceLoader'>
+<class 'lazy_imports_lite._loader.LazyLoader'>
+<class '_frozen_importlib_external.SourceFileLoader'>
+"""
+        ),
+        transformed_stderr=snapshot("<equal to normal>"),
+        normal_stdout=snapshot(
+            """\
+<class '_frozen_importlib_external.NamespaceLoader'>
+<class '_frozen_importlib_external.SourceFileLoader'>
+<class '_frozen_importlib_external.SourceFileLoader'>
+"""
+        ),
+        normal_stderr=snapshot(""),
+    )