[BUG] Fix validation registry discovery when loaded via zipimport

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`. Generated modules under `expressions/generated/` are now discovered whether the package lives on disk or inside a zip archive.

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

@@ -7,15 +7,16 @@
 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
-`expressions/generated/` subtree simply yields an empty registry --
+`expressions/generated/` subtree simply yields an empty registry, and
 the package still imports cleanly.
 """
 
 from __future__ import annotations
 
 import importlib
 import logging
-from pathlib import Path
+import zipfile
+from pathlib import Path, PurePosixPath
 
 from .check import ModelValidation
 
@@ -24,23 +25,62 @@
 _GENERATED_ROOT = "overture.schema.pyspark.expressions.generated"
 
 
+def _zip_boundary(root_path: str) -> tuple[str, str] | None:
+    """Split a namespace portion into `(zip file path, internal prefix)`.
+
+    Returns `None` when `root_path` is neither a real directory nor inside a
+    zip, such as an empty or nonexistent portion. When Glue loads the package
+    from a wheel via `--extra-py-files`, a namespace portion looks like
+    `.../some_pkg-1.0-py3-none-any.whl/a/b/c`, where the leading
+    `.../some_pkg....whl` segment is a real zip file even though the whole
+    path is not a directory on disk. `importlib.resources`'s `Traversable`
+    API is meant to cover this, but its `MultiplexedPath` (at least through
+    Python 3.10) raises `NotADirectoryError` the moment any namespace portion
+    is not a real directory, so it does not work here.
+    """
+    path = Path(root_path)
+    for parent in (path, *path.parents):
+        if parent.is_file() and zipfile.is_zipfile(parent):
+            return str(parent), path.relative_to(parent).as_posix()
+    return None
+
+
 def _iter_generated_module_names(root_paths: list[str]) -> list[str]:
-    """Return the dotted names of every generated module on disk.
+    """Return the dotted names of every generated module under `root_paths`.
 
     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. Each
+    namespace portion is therefore walked directly. A portion that is a real
+    directory is walked with `pathlib`; one that sits inside a wheel on
+    `sys.path`, as under Glue's `--extra-py-files`, is read from the archive
+    with `zipfile`.
     """
     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
+        if base.is_dir():
+            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]))
+            continue
+
+        boundary = _zip_boundary(root_path)
+        if boundary is None:
+            continue
+        zip_path, prefix = boundary
+        prefix = f"{prefix}/" if prefix else ""
+        with zipfile.ZipFile(zip_path) as archive:
+            for entry in archive.namelist():
+                if not entry.startswith(prefix) or not entry.endswith(".py"):
+                    continue
+                member = PurePosixPath(entry[len(prefix) :])
+                if member.name == "__init__.py":
+                    continue
+                dotted = member.with_suffix("").as_posix().replace("/", ".")
+                names.append(".".join([_GENERATED_ROOT, dotted]))
+    return sorted(names)
 
 
 def _walk() -> tuple[dict[str, ModelValidation], dict[str, dict[str, str]]]:

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

@@ -2,28 +2,35 @@
 
 The generated expression tree is PEP 420 (no `__init__.py`), so the
 registry must walk it as a namespace package. No other test exercises the
-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.
+real on-disk walk, since 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 zip branch of the walk (a namespace portion inside a wheel on
+`sys.path`, as on Glue) has no such incidental coverage either, so it is
+exercised directly here against a synthetic archive.
 """
 
 from __future__ import annotations
 
 import importlib
+import zipfile
 from pathlib import Path
 
 import pytest
 
-from overture.schema.pyspark._registry import REGISTRY
-
-_GENERATED_ROOT = "overture.schema.pyspark.expressions.generated"
+from overture.schema.pyspark._registry import (
+    _GENERATED_ROOT,
+    REGISTRY,
+    _iter_generated_module_names,
+)
 
 
 def _generated_leaf_count() -> int:
     """Count generated model modules on disk (excludes namespace dirs).
 
-    Returns 0 when the generated tree is absent -- mirroring the registry's
-    own `ImportError` handling -- so the test skips rather than errors.
+    Returns 0 when the generated tree is absent, matching the registry's own
+    `ImportError` handling, so the test skips cleanly.
     """
     try:
         root = importlib.import_module(_GENERATED_ROOT)
@@ -46,3 +53,33 @@ 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_zip(tmp_path: Path) -> None:
+    """The zip branch lists generated modules inside a wheel on `sys.path`.
+
+    Mirrors how Glue loads the package via `--extra-py-files`: the namespace
+    portion is `<wheel>/overture/schema/pyspark/expressions/generated`, and
+    its leading `<wheel>` segment is a real zip file. Namespace `__init__.py`
+    markers and members outside the generated prefix are excluded, and the
+    two feature modules come back as dotted names.
+    """
+    prefix = "overture/schema/pyspark/expressions/generated"
+    wheel = tmp_path / "overture_schema_pyspark-0.0.0-py3-none-any.whl"
+    with zipfile.ZipFile(wheel, "w") as archive:
+        archive.writestr(f"{prefix}/overture/schema/buildings/building.py", "\n")
+        archive.writestr(f"{prefix}/overture/schema/base/water.py", "\n")
+        archive.writestr(f"{prefix}/overture/schema/base/__init__.py", "\n")
+        archive.writestr("overture/schema/pyspark/check.py", "\n")
+
+    names = _iter_generated_module_names([f"{wheel}/{prefix}"])
+
+    assert names == [
+        f"{_GENERATED_ROOT}.overture.schema.base.water",
+        f"{_GENERATED_ROOT}.overture.schema.buildings.building",
+    ]
+
+
+def test_iter_generated_module_names_ignores_nonexistent_path(tmp_path: Path) -> None:
+    """A portion that is neither a directory nor inside a zip yields nothing."""
+    assert _iter_generated_module_names([str(tmp_path / "missing")]) == []