[BUG] Discover validation registry through importlib.resources

packages/overture-schema-pyspark/changelog.d/661.bugfix.md

@@ -0,0 +1 @@
+Fixed the validation registry coming back empty when `overture-schema-pyspark` is loaded from a wheel on `sys.path` (zipimport) rather than installed to a real directory, as happens on AWS Glue via `--extra-py-files`. The generated tree is now read through `importlib.resources`, which resolves a namespace portion inside an archive as well as one on disk.

packages/overture-schema-pyspark/pyproject.toml

@@ -5,6 +5,9 @@ build-backend = "uv_build"
 [project]
 dependencies = [
     "click>=8.0",
+    # 6.2 is the first release whose `files()` resolves a namespace portion
+    # inside a zip; the stdlib gained the same fix in 3.13.
+    "importlib-resources>=6.2; python_version < '3.13'",
     "overture-schema-system>=0.1.1",
     "packaging>=22",
 ]

packages/overture-schema-pyspark/src/overture/schema/pyspark/_registry.py

@@ -4,18 +4,38 @@
 namespace and collecting every module that exposes the
 codegen-emitted `ENTRY_POINT` and `MODEL_VALIDATION` constants.
 
-The generated tree on disk is the runtime source of truth: the
-registry contains exactly what was generated, regardless of which
-theme packages are installed alongside the pyspark package. A missing
+The generated tree is the runtime source of truth: the registry
+contains exactly what was generated, regardless of which theme
+packages are installed alongside the pyspark package. A missing
 `expressions/generated/` subtree simply yields an empty registry --
 the package still imports cleanly.
+
+The tree is read through `importlib.resources`, which resolves a
+namespace portion whether it is a directory on disk or a member of an
+archive. Any wheel left unextracted on `sys.path` is zipimported and
+`pathlib` cannot traverse into it -- Spark ships wheels that way with
+`--py-files`, as does AWS Glue with `--extra-py-files`.
 """
 
 from __future__ import annotations
 
 import importlib
 import logging
-from pathlib import Path
+import sys
+from collections.abc import Iterator
+
+if sys.version_info >= (3, 13):
+    from importlib.resources import files
+    from importlib.resources.abc import Traversable
+else:
+    # `importlib.resources.files` raises `NotADirectoryError` for a namespace
+    # package with any non-directory portion through Python 3.12; the backport
+    # carries the 3.13 fix. Runtimes known to sit below it, where this is always
+    # the branch taken: AWS Glue 4.0 (Python 3.10) and Glue 5.0 (3.11). Other
+    # runtimes that ship wheels unextracted belong on that list -- if you hit
+    # this somewhere else, please add what you find.
+    from importlib_resources import files
+    from importlib_resources.abc import Traversable
 
 from .check import ModelValidation
 
@@ -24,23 +44,28 @@
 _GENERATED_ROOT = "overture.schema.pyspark.expressions.generated"
 
 
-def _iter_generated_module_names(root_paths: list[str]) -> list[str]:
-    """Return the dotted names of every generated module on disk.
+def _iter_generated_module_names(root: str = _GENERATED_ROOT) -> list[str]:
+    """Return the dotted names of every generated module under `root`.
 
     The generated tree is PEP 420 (no `__init__.py`), so its subdirectories
-    are namespace packages. `pkgutil.walk_packages` skips those, so the tree
-    is walked as files instead: every `.py` under the namespace roots, keyed
-    to a dotted name relative to `_GENERATED_ROOT`.
+    are namespace packages, which `pkgutil.walk_packages` skips. It is walked
+    as resources instead: every `.py` below `root`, keyed to a dotted name.
+    `files` multiplexes every portion of the namespace, so a tree assembled
+    from more than one distribution is walked whole.
     """
-    names: list[str] = []
-    for root_path in root_paths:
-        base = Path(root_path)
-        for path in sorted(base.rglob("*.py")):
-            if path.name == "__init__.py":
-                continue
-            relative = path.relative_to(base).with_suffix("")
-            names.append(".".join([_GENERATED_ROOT, *relative.parts]))
-    return names
+
+    def walk(node: Traversable, prefix: tuple[str, ...]) -> Iterator[str]:
+        for child in node.iterdir():
+            if child.is_dir():
+                yield from walk(child, (*prefix, child.name))
+            elif child.name.endswith(".py") and child.name != "__init__.py":
+                yield ".".join([root, *prefix, child.name[: -len(".py")]])
+
+    try:
+        anchor = files(root)
+    except ModuleNotFoundError:
+        return []
+    return sorted(walk(anchor, ()))
 
 
 def _walk() -> tuple[dict[str, ModelValidation], dict[str, dict[str, str]]]:
@@ -61,12 +86,7 @@ def _walk() -> tuple[dict[str, ModelValidation], dict[str, dict[str, str]]]:
     registry: dict[str, ModelValidation] = {}
     partition_map: dict[str, dict[str, str]] = {}
 
-    try:
-        root = importlib.import_module(_GENERATED_ROOT)
-    except ImportError:
-        return registry, partition_map
-
-    for name in _iter_generated_module_names(list(root.__path__)):
+    for name in _iter_generated_module_names():
         module = importlib.import_module(name)
         entry_point = getattr(module, "ENTRY_POINT", None)
         validation = getattr(module, "MODEL_VALIDATION", None)

packages/overture-schema-pyspark/tests/test_registry.py

@@ -5,19 +5,86 @@
 real on-disk walk -- conformance tests import expression modules directly
 and `test_validate.py` registers models through a test shim -- so an empty
 registry would otherwise pass the suite unnoticed.
+
+The walk has to work for a namespace portion inside a zip as well as one on
+disk, because any wheel left unextracted on `sys.path` is zipimported --
+Spark's `--py-files` and Glue's `--extra-py-files` both ship wheels that way.
+Both shapes are exercised here against a synthetic package put on `sys.path`,
+so the portion strings under test come from real import machinery rather than
+being hand-written.
 """
 
 from __future__ import annotations
 
 import importlib
+import sys
+import zipfile
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
 from pathlib import Path
 
 import pytest
 
-from overture.schema.pyspark._registry import REGISTRY
+from overture.schema.pyspark._registry import REGISTRY, _iter_generated_module_names
 
 _GENERATED_ROOT = "overture.schema.pyspark.expressions.generated"
 
+_PROBE_TOP = "zipimport_probe"
+_PROBE_ROOT = f"{_PROBE_TOP}.expressions.generated"
+
+
+@contextmanager
+def _on_sys_path(entry: Path) -> Iterator[None]:
+    """Put `entry` on `sys.path` and unimport the probe package after."""
+    sys.path.insert(0, str(entry))
+    importlib.invalidate_caches()
+    try:
+        yield
+    finally:
+        sys.path.remove(str(entry))
+        for name in [n for n in sys.modules if n.split(".")[0] == _PROBE_TOP]:
+            del sys.modules[name]
+        importlib.invalidate_caches()
+
+
+def _probe_members(relative_paths: Sequence[str]) -> list[str]:
+    """Return archive member names for a probe tree holding `relative_paths`."""
+    return [f"{_PROBE_ROOT.replace('.', '/')}/{p}" for p in relative_paths]
+
+
+def _write_probe_wheel(wheel: Path, relative_paths: Sequence[str]) -> None:
+    """Write a wheel-shaped zip holding a PEP 420 probe tree.
+
+    Directory entries are written explicitly. `zipimport` recognises a
+    namespace portion inside an archive only when the archive carries a
+    directory entry for it, and that is what `uv_build` emits, so omitting
+    them here would test a shape no real wheel has.
+    """
+    members = _probe_members(relative_paths)
+    directories = sorted(
+        {
+            f"{parent}/"
+            for member in members
+            for parent in (str(p) for p in Path(member).parents)
+            if parent != "."
+        }
+    )
+    with zipfile.ZipFile(wheel, "w") as archive:
+        for directory in directories:
+            info = zipfile.ZipInfo(directory)
+            info.external_attr = (0o40755 << 16) | 0x10
+            archive.writestr(info, b"")
+        for member in members:
+            archive.writestr(member, "")
+
+
+def _write_probe_directory(root: Path, relative_paths: Sequence[str]) -> None:
+    """Write a PEP 420 probe tree as real directories under `root`."""
+    for member in _probe_members(relative_paths):
+        path = root / member
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text("")
+
 
 def _generated_leaf_count() -> int:
     """Count generated model modules on disk (excludes namespace dirs).
@@ -46,3 +113,38 @@ def test_registry_discovers_generated_models() -> None:
         key for key in REGISTRY if ":" in key and key.startswith("overture.schema.")
     ]
     assert generated_entries, "registry found no generated feature modules on disk"
+
+
+def test_iter_generated_module_names_reads_a_zipimported_tree(tmp_path: Path) -> None:
+    """Modules inside a wheel on `sys.path` are discovered.
+
+    The wheel is never unpacked, so the namespace portion's `__path__` points
+    inside the archive.
+    """
+    wheel = tmp_path / "zipimport_probe-0.0.0-py3-none-any.whl"
+    _write_probe_wheel(wheel, ["schema/base/water.py", "schema/buildings/building.py"])
+
+    with _on_sys_path(wheel):
+        names = _iter_generated_module_names(_PROBE_ROOT)
+
+    assert names == [
+        f"{_PROBE_ROOT}.schema.base.water",
+        f"{_PROBE_ROOT}.schema.buildings.building",
+    ]
+
+
+def test_iter_generated_module_names_reads_a_directory_tree(tmp_path: Path) -> None:
+    """Modules in a real directory are discovered, and `__init__.py` is skipped."""
+    _write_probe_directory(
+        tmp_path, ["schema/base/water.py", "schema/base/__init__.py"]
+    )
+
+    with _on_sys_path(tmp_path):
+        names = _iter_generated_module_names(_PROBE_ROOT)
+
+    assert names == [f"{_PROBE_ROOT}.schema.base.water"]
+
+
+def test_iter_generated_module_names_without_a_tree_is_empty() -> None:
+    """An absent generated tree yields no names rather than raising."""
+    assert _iter_generated_module_names(f"{_PROBE_ROOT}.absent") == []