diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index 7eae5a0..c7e9e11 100644 --- a/src/substrait/extension_registry/registry.py +++ b/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) diff --git a/tests/extension_registry/test_collector.py b/tests/extension_registry/test_collector.py index d36f10f..8326d78 100644 --- a/tests/extension_registry/test_collector.py +++ b/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 diff --git a/tests/extension_registry/test_default_extensions.py b/tests/extension_registry/test_default_extensions.py new file mode 100644 index 0000000..213482b --- /dev/null +++ b/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)