diff --git a/effectful/handlers/llm/harness/serialization.py b/effectful/handlers/llm/harness/serialization.py index 713de936f..abfd8553d 100644 --- a/effectful/handlers/llm/harness/serialization.py +++ b/effectful/handlers/llm/harness/serialization.py @@ -7,7 +7,6 @@ import abc import base64 import collections.abc -import contextvars import dataclasses import functools import inspect @@ -39,7 +38,6 @@ GenericAlias, TypeEvaluator, UnionType, - canonicalize, nested_type, ) from effectful.ops.types import Operation, Term @@ -352,74 +350,6 @@ class _NameAndTool(typing.NamedTuple): tool: Tool -# TODO move upstream to unification.py -@nested_type.register -def _nested_type_alias(ty: typing.TypeAliasType): - return nested_type(ty.__value__) - - -# TODO move upstream to unification.py -def _expand_alias( - evaluator: TypeEvaluator, typ: typing.TypeAliasType, value: typing.Any -): - """Evaluate what ``typ`` aliases, unless ``typ`` is already being expanded.""" - if not hasattr(evaluator, "_expanding_aliases"): - setattr(evaluator, "_expanding_aliases", set()) - seen: set[typing.Any] = getattr(evaluator, "_expanding_aliases") - if typ in seen: - return typ - seen.add(typ) - try: - return evaluator.evaluate(value) - finally: - seen.discard(typ) - - -# TODO move upstream to unification.py -@TypeEvaluator.evaluate.register # type: ignore[attr-defined] -def _evaluate_type_alias(self, typ: typing.TypeAliasType): - return _expand_alias(self, typ, typ.__value__) - - -# TODO move upstream to unification.py -@TypeEvaluator.evaluate.register # type: ignore[attr-defined] -def _evaluate_generic_alias(self, typ: GenericAlias): - origin, args = typing.get_origin(typ), typing.get_args(typ) - if isinstance(origin, typing.TypeAliasType): - return _expand_alias(self, typ, origin.__value__[args]) - return origin[self.evaluate(args)] # type: ignore[index] - - -# TODO move upstream to unification.py -_CANONICALIZING: contextvars.ContextVar[frozenset] = contextvars.ContextVar( - "_CANONICALIZING", default=frozenset() -) - - -# TODO move upstream to unification.py -@dataclasses.dataclass(frozen=True) -class _SelfReferentialAlias(Exception): - typ: typing.TypeAliasType - - -# TODO move upstream to unification.py -@canonicalize.register -def _canonicalize_type_alias(typ: typing.TypeAliasType) -> typing.Any: - """Canonicalize what the alias names, unless it names itself.""" - seen = _CANONICALIZING.get() - if typ in seen: - raise _SelfReferentialAlias(typ) - token = _CANONICALIZING.set(seen | {typ}) - try: - return canonicalize(typ.__value__) - except _SelfReferentialAlias as e: - if e.typ is not typ: - raise # another alias's recursion; the frame expanding it will catch - return typ - finally: - _CANONICALIZING.reset(token) - - class TypeToPydanticType(TypeEvaluator): """Substitute custom types with their Pydantic Annotated equivalents. diff --git a/effectful/internals/unification.py b/effectful/internals/unification.py index 7e351bbc2..6de3e2545 100644 --- a/effectful/internals/unification.py +++ b/effectful/internals/unification.py @@ -60,6 +60,7 @@ import builtins import collections import collections.abc +import contextvars import functools import inspect import numbers @@ -129,9 +130,36 @@ def evaluate(self, typ) -> TypeExpressions: def _(self, typ: TypeConstant | TypeVariable): return typ + def _expand_alias(self, typ, value): + """Evaluate what ``typ`` aliases, unless ``typ`` is already being expanded. + + A self-referential alias (``type JSON = int | list[JSON]``) has no finite + expansion, so the occurrence that closes the loop is left as the alias + itself and whoever consumes the result resolves it. The set of aliases + currently being expanded lives on the evaluator, so a frozen dataclass + subclass cannot be used with an alias. + """ + if not hasattr(self, "_expanding_aliases"): + setattr(self, "_expanding_aliases", set()) + seen: set = getattr(self, "_expanding_aliases") + # Key on the alias, not the subscripted form: an alias that reapplies + # itself at a *different* argument (``type T[X] = list[T[list[X]]]``) + # never repeats a subscription, so only the alias itself terminates it. + alias = typing.get_origin(typ) if isinstance(typ, GenericAlias) else typ + if alias in seen: + return typ + seen.add(alias) + try: + return self.evaluate(value) + finally: + seen.discard(alias) + @evaluate.register def _(self, typ: GenericAlias): origin, args = typing.get_origin(typ), typing.get_args(typ) + if isinstance(origin, typing.TypeAliasType): + # A subscripted generic alias: apply it, then evaluate what it names. + return self._expand_alias(typ, origin.__value__[args]) return origin[self.evaluate(args)] # type: ignore[index] @evaluate.register @@ -174,7 +202,7 @@ def _(self, typ: typing.NewType): @evaluate.register def _(self, typ: typing.TypeAliasType): - return self.evaluate(typ.__value__) + return self._expand_alias(typ, typ.__value__) @evaluate.register def _(self, typ: typing.ForwardRef): @@ -811,9 +839,39 @@ def _(typ: typing.NewType): return canonicalize(typ.__supertype__) +_CANONICALIZING: contextvars.ContextVar[frozenset] = contextvars.ContextVar( + "_CANONICALIZING", default=frozenset() +) + + +@dataclass(frozen=True) +class _SelfReferentialAlias(Exception): + typ: typing.TypeAliasType + + @canonicalize.register def _(typ: typing.TypeAliasType): - return canonicalize(typ.__value__) + """Canonicalize what the alias names, unless it names itself. + + A self-referential alias is its own canonical form: it has no finite + expansion, and keeping one expanded layer would not be idempotent -- each + further call would peel off another. So the recursive occurrence aborts the + whole expansion, back to the alias that started it. An alias reached *on the + way* re-raises, leaving its own frame to catch it, which is what makes + mutually recursive aliases terminate at the outermost of the two. + """ + seen = _CANONICALIZING.get() + if typ in seen: + raise _SelfReferentialAlias(typ) + token = _CANONICALIZING.set(seen | {typ}) + try: + return canonicalize(typ.__value__) + except _SelfReferentialAlias as e: + if e.typ is not typ: + raise # another alias's recursion; the frame expanding it will catch + return typ + finally: + _CANONICALIZING.reset(token) @canonicalize.register @@ -1114,6 +1172,17 @@ def _(value: str | bytes | range | None): return Box(type(value)) +@nested_type.register +def _(value: typing.TypeAliasType): + """An alias names a type, so it has the nested type of the type it names. + + Reached when an alias appears as a *value* rather than as an annotation -- + read out of a scope, say, or passed to an operation -- where reporting + `typing.TypeAliasType` would hide what the alias is for. + """ + return nested_type(value.__value__) + + _nested_type_dispatch = nested_type _nested_type_state = threading.local() _NESTED_TYPE_MAX_DEPTH = 5 diff --git a/tests/test_internals_unification.py b/tests/test_internals_unification.py index 35dfb77e7..1fc7d9bbe 100644 --- a/tests/test_internals_unification.py +++ b/tests/test_internals_unification.py @@ -36,6 +36,32 @@ W = typing.TypeVar("W") +# --- Type aliases --- +# +# An alias names a type; it does not make a new one, so every operation here must +# see through it to the answer it gives for the type it names. A *self-referential* +# alias has no finite expansion, so it is instead an opaque constant standing for +# itself -- the alternative is not a better answer but a `RecursionError`. +type _PairAlias = tuple[int, str] +type _PairsAlias = list[_PairAlias] +type _GenPairAlias[T] = tuple[T, T] +type _RecursiveAlias = int | list[_RecursiveAlias] +type _RecursiveGenAlias[T] = T | list[_RecursiveGenAlias[T]] + +# Mutually recursive: neither alias mentions itself directly, so expanding either +# terminates only if the guard survives a trip through the other. +type _MutualA = int | _MutualB +type _MutualB = str | list[_MutualA] + +# The pre-PEP-695 spelling, which is not a `TypeAliasType` at all -- it is the +# aliased type itself, bound to a name. +_LegacyPairAlias: typing.TypeAlias = tuple[int, str] # noqa: UP040 + + +class _Identity(TypeEvaluator): + """A `TypeEvaluator` that adds nothing, so it exercises the base traversal alone.""" + + @dataclasses.dataclass class _Substitute(TypeEvaluator): """ @@ -135,6 +161,11 @@ def _(self, typ: TypeVariable): (dict[K, collections.abc.Callable[[T], V]], {K, T, V}), # ParamSpec and TypeVarTuple (if needed later) # (collections.abc.Callable[typing.ParamSpec("P"), T], {T}), # Would need to handle ParamSpec + # Type aliases, which must not hide the variables of what they name + (_PairAlias, set()), + (list[_PairAlias], set()), + (_GenPairAlias[T], {T}), + (_RecursiveGenAlias[T], {T}), ], ) @pytest.mark.parametrize( @@ -2062,3 +2093,160 @@ def f[T](*args: T) -> T: typ = typeof(term) assert issubclass(typ, collections.abc.MutableMapping) assert not typing.is_typeddict(typ) + + +# ============================================================================ +# Type aliases (gh #766) +# +# A PEP 695 alias is transparent: it must be indistinguishable from the type it +# names, everywhere -- `canonicalize`, the `TypeEvaluator` traversal, `unify` and +# `nested_type` alike. The one case with no finite answer is a *self-referential* +# alias, which each operation must reach a fixpoint on rather than recurse into +# forever. +# ============================================================================ + +# (an alias, the type it names) +_ALIAS_CASES = [ + pytest.param(_PairAlias, tuple[int, str], id="tuple"), + pytest.param(_PairsAlias, list[tuple[int, str]], id="list-of-alias"), + pytest.param(_LegacyPairAlias, tuple[int, str], id="legacy"), + # The alias under a generic it did not itself introduce. + pytest.param(list[_PairAlias], list[tuple[int, str]], id="under-generic"), + pytest.param( + typing.Annotated[_PairAlias, "m"], + typing.Annotated[tuple[int, str], "m"], + id="annotated", + ), + # A generic alias, at the point it is applied. The alias loses its expansion + # by being *subscripted*, which is what the unsubscripted cases cannot catch. + pytest.param(_GenPairAlias[int], tuple[int, int], id="generic-subscripted"), +] + +# Aliases with no finite expansion. +_RECURSIVE_ALIAS_CASES = [ + pytest.param(_RecursiveAlias, id="plain"), + pytest.param(_RecursiveGenAlias[int], id="generic"), + pytest.param(_MutualA, id="mutual"), +] + + +@pytest.mark.parametrize("alias,target", _ALIAS_CASES) +def test_canonicalize_alias_is_transparent(alias, target): + assert canonicalize(alias) == canonicalize(target) + + +@pytest.mark.parametrize("alias,target", _ALIAS_CASES) +def test_type_evaluator_alias_is_transparent(alias, target): + """The base traversal expands an alias, including a subscripted generic one. + + Left unexpanded, a downstream consumer dispatching on the result sees a + `TypeAliasType` where a type was promised, and any handling registered for + what the alias names is silently skipped. + """ + assert _Identity().evaluate(alias) == _Identity().evaluate(target) + + +@pytest.mark.parametrize("alias,target", _ALIAS_CASES) +def test_nested_type_alias_matches_target(alias, target): + """An alias as a *value* reports what the type it names would report.""" + assert nested_type(alias) == nested_type(target) + + +@pytest.mark.parametrize("alias,target", _ALIAS_CASES) +def test_unify_alias_is_transparent(alias, target): + assert unify(alias, target) == {} + assert unify(target, alias) == {} + assert unify(T, alias) == {T: canonicalize(target)} + + +@pytest.mark.parametrize("alias", _RECURSIVE_ALIAS_CASES) +def test_canonicalize_recursive_alias_terminates(alias): + """A self-referential alias canonicalizes to itself, and stays there. + + Idempotency is the reason it collapses all the way back to the alias instead + of keeping one expanded layer: a layer per call is not a canonical form. + """ + assert canonicalize(alias) == alias + assert canonicalize(canonicalize(alias)) == canonicalize(alias) + + +def test_canonicalize_recursive_alias_under_a_generic(): + """The alias is opaque; the type carrying it is canonicalized as usual.""" + assert ( + canonicalize(list[_RecursiveAlias]) + == collections.abc.MutableSequence[_RecursiveAlias] + ) + + +@pytest.mark.parametrize( + "alias,expected", + [ + pytest.param(_RecursiveAlias, int | list[_RecursiveAlias], id="plain"), + pytest.param( + _RecursiveGenAlias[int], int | list[_RecursiveGenAlias[int]], id="generic" + ), + # Expansion stops only where it comes back around to the alias it + # started from, so `_MutualB` -- reached on the way -- is expanded too. + pytest.param(_MutualA, int | str | list[_MutualA], id="mutual"), + ], +) +def test_type_evaluator_recursive_alias_terminates(alias, expected): + """The traversal expands the body and leaves the recursive occurrence alone. + + Unlike `canonicalize`, the evaluator has no idempotency to preserve -- it + hands the result to a consumer (Pydantic, say) that resolves the remaining + alias itself -- so it keeps the expanded layer rather than collapsing back. + """ + assert _Identity().evaluate(alias) == expected + + +def test_unify_recursive_alias_is_opaque(): + """A recursive alias unifies as a constant: with itself, or with a variable. + + ``unify`` is not alias-aware; it canonicalizes its arguments first, so this + is what the canonicalization fixpoint buys it. Unifying the alias against + its own expansion is beyond that and fails -- but it *fails*, with the + ordinary error, rather than recursing until the stack runs out. + """ + assert unify(_RecursiveAlias, _RecursiveAlias) == {} + assert unify(T, _RecursiveAlias) == {T: _RecursiveAlias} + with pytest.raises(TypeError): + unify(_RecursiveAlias, int | list[_RecursiveAlias]) + + +def test_nested_type_recursive_alias_terminates(): + assert nested_type(_RecursiveAlias) == nested_type(int | list[_RecursiveAlias]) + + +def test_type_evaluator_alias_guard_does_not_leak(): + """The in-progress alias is released once its expansion is done. + + A guard that is set but never unset would leave every later alias + unexpanded -- and would do so only after a recursive one had been seen, + which no single-alias test would catch. + """ + evaluator = _Identity() + first = evaluator.evaluate(_RecursiveAlias) + assert evaluator.evaluate(_RecursiveAlias) == first + assert evaluator.evaluate(_PairAlias) == tuple[int, str] + assert evaluator.evaluate(list[_RecursiveAlias]) == list[_RecursiveAlias.__value__] + + +def test_canonicalize_alias_guard_does_not_leak(): + first = canonicalize(_RecursiveAlias) + assert canonicalize(_RecursiveAlias) == first + assert canonicalize(dict[str, _PairAlias]) == canonicalize( + dict[str, tuple[int, str]] + ) + + +def test_substitute_alias_agrees_with_evaluator(): + """The evaluator's alias expansion agrees with ground-truth substitution. + + Only up to canonicalization: ``substitute`` reapplies the alias + (``_GenPairAlias[int]``) where the evaluator expands it (``tuple[int, int]``), + and canonicalizing is what makes those the same type expression. + """ + assert canonicalize( + _Substitute.substitute(_GenPairAlias[T], {T: int}) + ) == canonicalize(substitute(_GenPairAlias[T], {T: int}))