fix(extension_registry): load default extensions via iterdir()

src/substrait/extension_registry/registry.py

@@ -6,6 +6,11 @@
 from pathlib import Path
 from typing import Optional, Union
 
+try:  # Python 3.11+
+    from importlib.resources.abc import Traversable
+except ImportError:  # Python 3.10
+    from importlib.abc import Traversable
+
 import yaml
 from substrait.type_pb2 import Type
 
@@ -34,21 +39,35 @@ def __init__(self, load_default_extensions=True) -> None:
         # extension relation's output schema can be derived during inference.
         self._extension_relations: dict = {}
         if load_default_extensions:
-            for fpath in importlib_files("substrait_extensions.extensions").glob(  # type: ignore
-                "functions*.yaml"
-            ):
+            # Traversable has no portable glob() (MultiplexedPath lacks it), so
+            # filter iterdir() by name instead.
+            ext_dir = importlib_files("substrait_extensions.extensions")
+            yamls = [
+                f
+                for f in ext_dir.iterdir()
+                if f.name.startswith("functions") and f.name.endswith(".yaml")
+            ]
+            if not yamls:
+                raise RuntimeError(
+                    f"no functions*.yaml found in {ext_dir!r}; "
+                    "is the substrait_extensions package data installed?"
+                )
+            for fpath in yamls:
                 self.register_extension_yaml(fpath)
 
     def register_extension_yaml(
         self,
-        fname: Union[str, Path],
+        fname: Union[str, Path, Traversable],
     ) -> None:
         """Register extensions from a YAML file.
         Args:
-            fname: Path to the YAML file
+            fname: Path to the YAML file, or any ``Traversable`` such as an entry
+                from ``importlib.resources.files``
         """
-        fname = Path(fname)
-        with open(fname) as f:  # type: ignore
+        if not hasattr(fname, "open"):
+            fname = Path(fname)
+        # Binary mode lets PyYAML detect the encoding instead of using the locale's.
+        with fname.open("rb") as f:
             extension_definitions = yaml.safe_load(f)
         self.register_extension_dict(extension_definitions)
 

tests/extension_registry/test_collector.py

@@ -66,13 +66,8 @@ def arithmetic_only_registry():
     """
     registry = ExtensionRegistry(load_default_extensions=False)
     registry.register_extension_yaml(
-        next(
-            iter(
-                importlib_resources.files("substrait_extensions.extensions").glob(
-                    "functions_arithmetic.yaml"
-                )
-            )
-        )
+        importlib_resources.files("substrait_extensions.extensions")
+        / "functions_arithmetic.yaml"
     )
     return registry
 

tests/extension_registry/test_default_extensions.py

@@ -0,0 +1,77 @@
+"""Regression tests for loading default extensions off a non-``Path`` resource.
+
+``importlib.resources.files`` returns a ``Traversable``, not necessarily a
+filesystem ``Path``. The registry must load the bundled extension YAMLs
+regardless of which concrete ``Traversable`` the resource reader hands back:
+
+* ``MultiplexedPath`` -- returned for a *namespace* package (no ``__init__.py``).
+  It implements ``iterdir``/``open``/``joinpath`` but *not* ``glob``, so the old
+  ``files(...).glob("functions*.yaml")`` failed with
+  ``AttributeError: 'MultiplexedPath' object has no attribute 'glob'`` (observed
+  under the pure-Python WASI guest interpreter).
+* ``zipfile.Path`` -- returned for a zip-imported package. It is not an
+  ``os.PathLike``, so loading failed: before Python 3.12 ``zipfile.Path`` has no
+  ``glob`` (the old code raised ``AttributeError``), and feeding it to
+  ``Path(fname)`` raised ``TypeError``. ``register_extension_yaml`` now opens the
+  ``Traversable`` directly.
+"""
+
+import zipfile
+from importlib.resources import files
+
+import yaml
+
+import substrait.extension_registry.registry as registry_module
+from substrait.extension_registry import ExtensionRegistry
+
+try:  # Python 3.11+
+    from importlib.resources.readers import MultiplexedPath
+except ModuleNotFoundError:  # Python 3.10
+    from importlib.readers import MultiplexedPath
+
+
+def _load_defaults_from(monkeypatch, traversable) -> None:
+    """Build a default registry with ``files()`` resolving to ``traversable``."""
+    calls = []
+    monkeypatch.setattr(
+        registry_module,
+        "importlib_files",
+        lambda pkg: calls.append(pkg) or traversable,
+    )
+    reg = ExtensionRegistry(load_default_extensions=True)
+    assert calls == ["substrait_extensions.extensions"]
+    expected = {
+        yaml.safe_load(f.read_bytes())["urn"]
+        for f in files("substrait_extensions.extensions").iterdir()
+        if f.name.startswith("functions") and f.name.endswith(".yaml")
+    }
+    assert set(reg.urns()) == expected
+
+
+def test_load_default_extensions_via_multiplexed_path(monkeypatch):
+    # Wrap the real extensions directory in a MultiplexedPath so ``files()``
+    # yields exactly what it would for a namespace package. Guard the premise:
+    # MultiplexedPath must not expose ``glob`` (otherwise this test proves
+    # nothing).
+    ext_dir = files("substrait_extensions.extensions")
+    multiplexed = MultiplexedPath(ext_dir)
+    assert not hasattr(multiplexed, "glob")
+
+    # The old files(...).glob(...) failed with AttributeError here.
+    _load_defaults_from(monkeypatch, multiplexed)
+
+
+def test_load_default_extensions_via_zipfile_path(monkeypatch, tmp_path):
+    # Mirror a zip-imported install: copy the real YAMLs into a zip and resolve
+    # ``files()`` to a zipfile.Path over the archive. The entries are not
+    # os.PathLike, so the old loader failed on them -- AttributeError on the
+    # missing ``glob`` before 3.12, TypeError in ``Path(fname)`` otherwise.
+    ext_dir = files("substrait_extensions.extensions")
+    archive = tmp_path / "substrait_extensions.zip"
+    with zipfile.ZipFile(archive, "w") as zf:
+        for child in ext_dir.iterdir():
+            if child.name.endswith(".yaml"):
+                zf.writestr(f"extensions/{child.name}", child.read_bytes())
+
+    zip_dir = zipfile.Path(archive, "extensions/")
+    _load_defaults_from(monkeypatch, zip_dir)