diff --git a/src/sciline/_unification.py b/src/sciline/_unification.py new file mode 100644 index 00000000..11337d07 --- /dev/null +++ b/src/sciline/_unification.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2025 Scipp contributors (https://github.com/scipp) +"""Unification of generic type patterns with concrete keys. + +Generic providers whose type variables lack constraints are not expanded +eagerly. They are kept as templates and instantiated on demand by unifying +their argument and return-type patterns with the concrete keys that appear in +the pipeline (parameters, requested targets, mapped keys). +""" + +from __future__ import annotations + +import itertools +from collections.abc import Generator, Iterable +from typing import TYPE_CHECKING, TypeVar, get_args, get_origin + +if TYPE_CHECKING: + from ._provider import Provider + from .typing import Key + + +def find_all_typevars(t: type | TypeVar) -> set[TypeVar]: + """Returns the set of all TypeVars in a type expression.""" + if isinstance(t, TypeVar): + return {t} + if params := getattr(t, '__parameters__', ()): + return set(params) + return set(itertools.chain(*map(find_all_typevars, get_args(t)))) + + +def unify(pattern: Key | TypeVar, concrete: Key, bound: dict[TypeVar, Key]) -> bool: + """Match ``concrete`` against ``pattern``, extending ``bound`` in place. + + Returns True on success. ``bound`` may contain partial bindings on failure + and must be discarded by the caller in that case. + """ + if isinstance(pattern, TypeVar): + if pattern.__constraints__ and concrete not in pattern.__constraints__: + return False + if pattern in bound: + return bound[pattern] == concrete + bound[pattern] = concrete + return True + if (origin := get_origin(pattern)) is None: + return pattern == concrete + if get_origin(concrete) != origin: + return False + pattern_args = get_args(pattern) + concrete_args = get_args(concrete) + if len(pattern_args) != len(concrete_args): + return False + return all( + unify(p, c, bound) for p, c in zip(pattern_args, concrete_args, strict=True) + ) + + +def match_return(template: Provider, key: Key) -> Provider | None: + """Instantiate a generic provider if its return type unifies with ``key``.""" + bound: dict[TypeVar, Key] = {} + if unify(template.deduce_key(), key, bound): + return template.bind_type_vars(bound) + return None + + +def forward_bindings( + template: Provider, known_keys: Iterable[Key] +) -> Generator[tuple[dict[TypeVar, Key], frozenset[Key]], None, None]: + """Enumerate complete bindings of a template's TypeVars from known keys. + + Each generic argument of the template is unified with each known key; + consistent combinations of the resulting bindings that bind all type + variables of the template are yielded, together with the set of known keys + that produced them. Arguments that match no known key are left unmatched, + i.e., a binding is complete as long as the *other* arguments determine all + type variables. + """ + typevars = find_all_typevars(template.deduce_key()) + patterns = [p for p in template.arg_spec.keys() if find_all_typevars(p)] + options = [] + for pattern in patterns: + matches: list[tuple[dict[TypeVar, Key], Key | None]] = [({}, None)] + for key in known_keys: + bound: dict[TypeVar, Key] = {} + if unify(pattern, key, bound): + matches.append((bound, key)) + options.append(matches) + seen = set() + for combo in itertools.product(*options): + merged: dict[TypeVar, Key] = {} + if not all(_merge(merged, bound) for bound, _ in combo): + continue + if set(merged) != typevars: + continue + used = frozenset(key for _, key in combo if key is not None) + if (fingerprint := (frozenset(merged.items()), used)) not in seen: + seen.add(fingerprint) + yield merged, used + + +def _merge(target: dict[TypeVar, Key], bound: dict[TypeVar, Key]) -> bool: + for tv, key in bound.items(): + if target.setdefault(tv, key) != key: + return False + return True diff --git a/src/sciline/data_graph.py b/src/sciline/data_graph.py index bd754e03..a883e1ac 100644 --- a/src/sciline/data_graph.py +++ b/src/sciline/data_graph.py @@ -6,13 +6,20 @@ import itertools from collections.abc import Callable, Generator, Iterable, Mapping from types import NoneType -from typing import TYPE_CHECKING, Any, TypeVar, get_args +from typing import TYPE_CHECKING, Any, TypeVar, get_origin import cyclebane as cb import networkx as nx from cyclebane.node_values import IndexName, IndexValue -from ._provider import ArgSpec, Provider, ToProvider, _bind_free_typevars +from ._provider import ( + ArgSpec, + Provider, + ToProvider, + UnboundTypeVar, + _bind_free_typevars, +) +from ._unification import find_all_typevars, forward_bindings, match_return, unify from ._utils import key_full_qualname from .handler import ErrorHandler, HandleAsBuildTimeException from .typing import Graph, Key @@ -28,15 +35,6 @@ def _as_graph(key: Key, value: Any) -> cb.Graph: return cb.Graph(graph) -def _find_all_typevars(t: type | TypeVar) -> set[TypeVar]: - """Returns the set of all TypeVars in a type expression.""" - if isinstance(t, TypeVar): - return {t} - if params := getattr(t, '__parameters__', ()): - return set(params) - return set(itertools.chain(*map(_find_all_typevars, get_args(t)))) - - def _get_typevar_constraints( t: TypeVar, over_constraints: dict[TypeVar, frozenset[Key]] ) -> frozenset[Key]: @@ -81,6 +79,8 @@ def _normalize_custom_constraints( T = TypeVar('T', bound='DataGraph') +_providing_attrs = frozenset(('value', 'provider', 'reduce')) + class DataGraph: def __init__( @@ -90,20 +90,26 @@ def __init__( constraints: Mapping[TypeVar, Iterable[Key]] | None = None, ) -> None: self._constraints = _normalize_custom_constraints(constraints) + self._templates: list[Provider] = [] + self._template_values: dict[Key, Any] = {} self._cbgraph = cb.Graph(nx.DiGraph()) for provider in providers or []: self.insert(provider) - @classmethod - def _from_cyclebane(cls: type[T], graph: cb.Graph) -> T: - out = cls([]) + def _from_cyclebane(self: T, graph: cb.Graph) -> T: + out = type(self)([]) out._cbgraph = graph + out._constraints = self._constraints + out._templates = list(self._templates) + out._template_values = dict(self._template_values) return out + @property + def _has_templates(self) -> bool: + return bool(self._templates or self._template_values) + def copy(self: T) -> T: - cpy = self._from_cyclebane(self._cbgraph.copy()) - cpy._constraints = self._constraints - return cpy + return self._from_cyclebane(self._cbgraph.copy()) def __copy__(self: T) -> T: return self.copy() @@ -152,9 +158,12 @@ def insert(self, provider: ToProvider | Provider, /) -> None: if not isinstance(provider, Provider): provider = Provider.from_function(provider) return_type = provider.deduce_key() - if typevars := _find_all_typevars(return_type): - for bound in _mapping_to_constrained(typevars, self._constraints): - self.insert(provider.bind_type_vars(bound)) + if typevars := find_all_typevars(return_type): + if all(t.__constraints__ or t in self._constraints for t in typevars): + for bound in _mapping_to_constrained(typevars, self._constraints): + self.insert(provider.bind_type_vars(bound)) + else: + self._register_template(provider, typevars) return # Trigger UnboundTypeVar error if any input typevars are not bound provider = provider.bind_type_vars({}) @@ -162,6 +171,113 @@ def insert(self, provider: ToProvider | Provider, /) -> None: for dep in provider.arg_spec.keys(): self.underlying_graph.add_edge(dep, return_type, key=dep) + def _register_template(self, provider: Provider, typevars: set[TypeVar]) -> None: + """Store a generic provider for on-demand instantiation. + + Providers with unconstrained type variables cannot be expanded eagerly. + They are instantiated later by unifying their type patterns with the + concrete keys that appear in the graph, see :py:meth:`_instantiate_backward` + and :py:meth:`_instantiate_forward`. + """ + return_type = provider.deduce_key() + if isinstance(return_type, TypeVar): + raise ValueError( + f"Provider {provider} returns a bare unconstrained type variable " + f"{return_type!r}, which would match any requested key. Use a " + "generic class as return type or constrain the type variable." + ) + arg_typevars: set[TypeVar] = set() + for arg in provider.arg_spec.keys(): + arg_typevars |= find_all_typevars(arg) + if unbound := arg_typevars - typevars: + raise UnboundTypeVar( + f"Provider {provider} has type variables {unbound} in its " + "arguments that do not appear in its return type." + ) + # Mirror the replacement semantics of inserting a concrete provider twice. + self._templates = [t for t in self._templates if t != provider] + self._templates.append(provider) + + def _register_template_value(self, key: Key, value: Any) -> None: + """Store a value for a generic key, applied to all demanded specializations.""" + if get_origin(key) is None and (params := getattr(key, '__parameters__', ())): + # Normalize a bare generic class to a subscripted pattern. + key = key[params] # type: ignore[index] + # Mirror the replacement semantics of setting a concrete key twice. + self._template_values.pop(key, None) + self._template_values[key] = value + + def _match_template_value(self, key: Key) -> Any: + for pattern, value in reversed(self._template_values.items()): + bound: dict[TypeVar, Key] = {} + if unify(pattern, key, bound): + return value + return _no_value + + def _satisfied(self, key: Key) -> bool: + graph = self.underlying_graph + return key in graph and bool(graph.nodes[key].keys() & _providing_attrs) + + def _instantiate_backward(self, keys: Iterable[Key]) -> None: + """Instantiate templates for demanded keys and their dependencies.""" + stack = list(keys) + seen = set() + while stack: + key = stack.pop() + if key in seen: + continue + seen.add(key) + if not self._satisfied(key): + if (value := self._match_template_value(key)) is not _no_value: + self[key] = value + else: + # Iterate in reverse so that the latest matching template wins, + # mirroring the replacement semantics of concrete providers. + for template in reversed(self._templates): + if (provider := match_return(template, key)) is not None: + self.insert(provider) + break + if key in self.underlying_graph: + stack.extend(self.underlying_graph.predecessors(key)) + + def _instantiate_forward(self, seeds: Iterable[Key] | None = None) -> None: + """Instantiate templates whose arguments unify with concrete keys. + + If ``seeds`` is given, only instantiations consuming a seed key or a + key derived from one are created; other keys may still contribute to + bindings. Otherwise all complete bindings from the present concrete + keys are instantiated. Runs to a fixed point since instantiated + providers introduce new keys. + """ + derived = None if seeds is None else set(seeds) + instantiated: set[Key] = set() + while True: + known = set(self.underlying_graph.nodes) + if derived is not None: + known |= derived + providers: dict[Key, Provider] = {} + for template in self._templates: + for bound, used in forward_bindings(template, known): + if derived is not None and not (used & derived): + continue + provider = template.bind_type_vars(bound) + providers[provider.deduce_key()] = provider + inserted = False + for key, provider in providers.items(): + if key not in instantiated and not self._satisfied(key): + self.insert(provider) + instantiated.add(key) + if derived is not None: + derived.add(key) + inserted = True + if not inserted: + for key in list(self.underlying_graph.nodes): + if not self._satisfied(key) and ( + (value := self._match_template_value(key)) is not _no_value + ): + self[key] = value + return + def __setitem__(self, key: Key, value: DataGraph | Any) -> None: """ Provide a concrete value for a type. @@ -176,9 +292,12 @@ def __setitem__(self, key: Key, value: DataGraph | Any) -> None: # This is a questionable approach: Using MyGeneric[T] as a key will actually # not pass mypy [valid-type] checks. What we do on our side is ok, but the # calling code is not. - if typevars := _find_all_typevars(key): - for bound in _mapping_to_constrained(typevars, self._constraints): - self[_bind_free_typevars(key, bound)] = value + if typevars := find_all_typevars(key): + if all(t.__constraints__ or t in self._constraints for t in typevars): + for bound in _mapping_to_constrained(typevars, self._constraints): + self[_bind_free_typevars(key, bound)] = value + else: + self._register_template_value(key, value) return # TODO If key is generic, should we support multi-sink case and update all? @@ -189,7 +308,11 @@ def __setitem__(self, key: Key, value: DataGraph | Any) -> None: def __getitem__(self: T, key: Key) -> T: """Return the subgraph that computes the given key.""" - return self._from_cyclebane(self._cbgraph[key]) + graph = self + if self._has_templates: + graph = self.copy() + graph._instantiate_backward((key,)) + return graph._from_cyclebane(graph._cbgraph[key]) def map(self: T, node_values: dict[Key, Any]) -> T: """Map the graph over given node values. @@ -207,7 +330,13 @@ def map(self: T, node_values: dict[Key, Any]) -> T: : A new graph with mapped nodes. """ - return self._from_cyclebane(self._cbgraph.map(node_values)) + graph = self + if self._has_templates: + # Mapping duplicates dependents of the mapped nodes, so generic + # providers must be instantiated first. + graph = self.copy() + graph._instantiate_forward(node_values.keys()) + return graph._from_cyclebane(graph._cbgraph.map(node_values)) def reduce(self: T, *, func: Callable[..., Any], **kwargs: Any) -> T: """Reduce the outputs of a mapped graph into a single value and provider. @@ -258,6 +387,9 @@ def visualize_data_graph(self, **kwargs: Any) -> graphviz.Digraph: def to_task_graph( data_graph: DataGraph, targets: tuple[Key, ...], handler: ErrorHandler | None = None ) -> Graph: + if data_graph._has_templates: + data_graph = data_graph.copy() + data_graph._instantiate_backward(targets) graph = data_graph.to_networkx() handler = handler or HandleAsBuildTimeException() ancestors = list(targets) diff --git a/src/sciline/pipeline.py b/src/sciline/pipeline.py index fc94ef7b..1924b186 100644 --- a/src/sciline/pipeline.py +++ b/src/sciline/pipeline.py @@ -269,7 +269,13 @@ def get( graph = to_task_graph(self, targets=targets, handler=handler) # type: ignore[arg-type] except UnsatisfiedRequirement as e: missing = e.args[1] - nx_graph = self.underlying_graph + source = self + if self._has_templates: + # Instantiate generic providers so that the error message reflects + # the graph that was actually built. + source = self.copy() + source._instantiate_backward(targets) # type: ignore[arg-type] + nx_graph = source.underlying_graph if missing in nx_graph: paths = _find_paths_to_targets(nx_graph, missing, targets) info = _format_paths_msg(nx_graph, paths) @@ -345,8 +351,14 @@ def _repr_html_(self) -> str: def output_keys(self) -> tuple[Key, ...]: """Returns the keys that are not inputs to any other providers.""" + graph = self + if self._has_templates: + # Instantiate generic providers derivable from the present concrete + # keys so that their outputs are included. + graph = self.copy() + graph._instantiate_forward() sink_nodes = [ - node for node, degree in self.underlying_graph.out_degree if degree == 0 + node for node, degree in graph.underlying_graph.out_degree if degree == 0 ] return tuple(sorted(sink_nodes, key=key_name)) diff --git a/tests/conftest.py b/tests/conftest.py index 6691fce4..a25c6d99 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,11 +1,17 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2023 Scipp contributors (https://github.com/scipp) +import sys from typing import Any import pytest from sciline.scheduler import DaskScheduler, NaiveScheduler, Scheduler +collect_ignore: list[str] = [] +if sys.version_info < (3, 12): + # Uses PEP 695 syntax, which would be a SyntaxError. + collect_ignore.append("pep695_test.py") + @pytest.fixture def naive_scheduler() -> NaiveScheduler: diff --git a/tests/pep695_test.py b/tests/pep695_test.py new file mode 100644 index 00000000..53652ff4 --- /dev/null +++ b/tests/pep695_test.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2025 Scipp contributors (https://github.com/scipp) +"""Tests for generic providers using PEP 695 syntax (Python >= 3.12). + +The type variables have no constraints; instantiations are inferred from the +concrete keys appearing in the pipeline. This file is excluded from collection +on Python < 3.12 via ``collect_ignore`` in ``conftest.py``. +""" + +import pytest + +import sciline as sl +from sciline.handler import UnsatisfiedRequirement + +type A = int +type B = int + + +class Raw[Run](float): ... + + +class Processed[Run](float): ... + + +class Reduced[Run](float): ... + + +def process[Run](x: Raw[Run]) -> Processed[Run]: + return Processed[Run](x * 2) + + +def reduce_run[Run](x: Processed[Run]) -> Reduced[Run]: + return Reduced[Run](x + 1) + + +def test_generic_provider_instantiated_from_param() -> None: + pl = sl.Pipeline([process], params={Raw[A]: Raw[A](1.5)}) + assert pl.compute(Processed[A]) == 3.0 + + +def test_generic_provider_instantiated_per_requested_key() -> None: + pl = sl.Pipeline([process], params={Raw[A]: Raw[A](1.0), Raw[B]: Raw[B](2.0)}) + assert pl.compute(Processed[A]) == 2.0 + assert pl.compute(Processed[B]) == 4.0 + + +def test_chain_of_generic_providers() -> None: + pl = sl.Pipeline([process, reduce_run], params={Raw[A]: Raw[A](2.0)}) + assert pl.compute(Reduced[A]) == 5.0 + + +def test_generic_source_provider_instantiated_from_target() -> None: + def make[Run]() -> Raw[Run]: + return Raw[Run](7.0) + + pl = sl.Pipeline([make, process]) + assert pl.compute(Processed[A]) == 14.0 + + +def test_concrete_provider_shadows_generic() -> None: + def special() -> Processed[A]: + return Processed[A](0.5) + + pl = sl.Pipeline([process, special], params={Raw[B]: Raw[B](1.0)}) + assert pl.compute(Processed[A]) == 0.5 + assert pl.compute(Processed[B]) == 2.0 + + +def test_param_shadows_generic_provider() -> None: + pl = sl.Pipeline( + [process], params={Raw[A]: Raw[A](1.0), Processed[A]: Processed[A](5.0)} + ) + assert pl.compute(Processed[A]) == 5.0 + + +def test_later_generic_provider_wins() -> None: + def process2[Run](x: Raw[Run]) -> Processed[Run]: + return Processed[Run](x * 10) + + pl = sl.Pipeline([process, process2], params={Raw[A]: Raw[A](1.0)}) + assert pl.compute(Processed[A]) == 10.0 + + +def test_multiple_typevars_bound_from_target() -> None: + class Combined[R1, R2](float): ... + + def combine[R1, R2](x: Raw[R1], y: Processed[R2]) -> Combined[R1, R2]: + return Combined[R1, R2](x + y) + + pl = sl.Pipeline( + [combine, process], params={Raw[A]: Raw[A](1.0), Raw[B]: Raw[B](2.0)} + ) + assert pl.compute(Combined[A, B]) == 5.0 + assert pl.compute(Combined[A, A]) == 3.0 + + +def test_generic_type_alias() -> None: + type RawImage[Run] = float + type CleanImage[Run] = float + + def clean[Run](x: RawImage[Run]) -> CleanImage[Run]: + return x + 1.0 + + pl = sl.Pipeline([clean], params={RawImage[A]: 1.0}) + assert pl.compute(CleanImage[A]) == 2.0 + + +def test_constrained_pep695_typevar_expanded_eagerly() -> None: + def process2[Run: (A, B)](x: Raw[Run]) -> Processed[Run]: + return Processed[Run](x * 3) + + pl = sl.Pipeline([process2], params={Raw[A]: Raw[A](1.0), Raw[B]: Raw[B](2.0)}) + assert pl.compute(Processed[A]) == 3.0 + assert pl.compute(Processed[B]) == 6.0 + + +def test_map_over_generic_pipeline() -> None: + pl = sl.Pipeline([process]) + result = ( + pl.map({Raw[A]: [Raw[A](1.0), Raw[A](2.0)]}) + .reduce(func=lambda *v: sum(v), name='total') + .compute('total') + ) + assert result == 6.0 + + +def test_getitem_returns_subgraph_with_instantiated_generics() -> None: + pl = sl.Pipeline([process, reduce_run], params={Raw[A]: Raw[A](2.0)}) + sub = pl[Processed[A]] + assert sub.compute(Processed[A]) == 4.0 + + +def test_missing_dependency_of_instantiated_generic_raises() -> None: + pl = sl.Pipeline([process]) + with pytest.raises(UnsatisfiedRequirement, match='Raw'): + pl.compute(Processed[A]) + + +def test_generic_param_applies_to_all_specializations() -> None: + pl = sl.Pipeline([process], params={Raw: Raw(1.2)}) + assert pl.compute(Processed[A]) == pytest.approx(2.4) + assert pl.compute(Processed[B]) == pytest.approx(2.4) + + +def test_concrete_param_shadows_generic_param() -> None: + pl = sl.Pipeline([process], params={Raw: Raw(1.0), Raw[A]: Raw[A](5.0)}) + assert pl.compute(Processed[A]) == 10.0 + assert pl.compute(Processed[B]) == 2.0 + + +def test_generic_param_used_by_mapped_pipeline() -> None: + def combine(x: Processed[A], y: Raw[B]) -> float: + return x + y + + pl = sl.Pipeline([process, combine], params={Raw: Raw(1.0)}) + result = ( + pl.map({Raw[A]: [Raw[A](1.0), Raw[A](2.0)]}) + .reduce(func=lambda *v: sum(v), name='total') + .compute('total') + ) + assert result == 8.0 + + +def test_output_keys_include_derivable_generic_outputs() -> None: + pl = sl.Pipeline([process, reduce_run], params={Raw[A]: Raw[A](1.0)}) + assert Reduced[A] in pl.output_keys() + assert Reduced[B] not in pl.output_keys() diff --git a/tests/pipeline_test.py b/tests/pipeline_test.py index 00bd72b2..e81e5fbb 100644 --- a/tests/pipeline_test.py +++ b/tests/pipeline_test.py @@ -1513,7 +1513,7 @@ def f(x: A[T]) -> B[T]: pipeline.get(B[int]) -def test_type_vars_must_be_constrained() -> None: +def test_unconstrained_type_vars_are_instantiated_on_demand() -> None: T = TypeVar('T') @dataclass @@ -1527,8 +1527,10 @@ class B(Generic[T]): def foo(x: A[T]) -> B[T]: return B[T](x.v) - with pytest.raises(ValueError, match="no constraint"): - sl.Pipeline([foo]) + pipeline = sl.Pipeline([foo], params={A[int]: A[int](3)}) + assert pipeline.compute(B[int]) == B[int](3) + with pytest.raises(sl.handler.UnsatisfiedRequirement): + pipeline.get(B[float]) def test_custom_constraint_is_sufficient() -> None: