From 6e5de77995f3e9c86cee4ffcd8741704ddd053fa Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Wed, 30 Sep 2026 13:44:40 -0700 Subject: [PATCH] Add `Cond`, `UnpackMap`, `GetAssociated` to represent our syntax The plan here is that we will start moving to an AST-first setup (based on eventually PEP 849) in which we get an AST and translate that into `typing` objects. Currently, UnpackMap is set up like: `UnpackMap[Body, T, Target]` is equivalent to `*[Body for T in Iter[Target]]`, and `T` needs to be an already declared `TypeVar`. (An optional 4th argument allows putting an `if`.) An alternative would be to have it be `UnpackMap[Body, lambda T: Target]`. Thoughts? That's already kind of how we write `GenericCallable`, though I'm not sure that `GenericCallable` is really done in an acceptable way right now... --- tests/test_type_eval.py | 235 ++++++++++++++++++++++++++- typemap/type_eval/_eval_operators.py | 73 +++++++++ typemap/type_eval/_eval_typing.py | 13 +- typemap/typing.py | 60 +++++++ 4 files changed, 376 insertions(+), 5 deletions(-) diff --git a/tests/test_type_eval.py b/tests/test_type_eval.py index 6c32413b..8e11f82b 100644 --- a/tests/test_type_eval.py +++ b/tests/test_type_eval.py @@ -27,7 +27,17 @@ from typemap.type_eval import _ensure_context, eval_typing from typemap.type_eval._eval_operators import TypeMapError -from typemap.typing import _BoolLiteral +from typemap.typing import ( + _All, + _And, + _Any, + _BoolLiteral, + _Cond, + _GetAssociated, + _Not, + _Or, + _UnpackMap, +) from typemap_extensions import ( Attrs, @@ -1670,6 +1680,229 @@ def test_eval_bool_05(): assert d == Literal[False] +type OnlyIntToSetCond[T] = _Cond[IsAssignable[T, int], set[T], T] +type CondIsIntLiteral[T] = _Cond[IsIntBool[T], Literal[True], Literal[False]] + + +def test_eval_cond_01(): + assert eval_typing(_Cond[Literal[True], int, str]) is int + assert eval_typing(_Cond[Literal[False], int, str]) is str + assert eval_typing(_Cond[Never, int, str]) is str + assert eval_typing(_Cond[int, int, str]) is str + assert eval_typing(_Cond[Literal[False] | Literal[True], int, str]) is int + + assert eval_typing(OnlyIntToSetCond[int]) == set[int] + assert eval_typing(OnlyIntToSetCond[str]) is str + + assert eval_typing(CondIsIntLiteral[int]) == Literal[True] + assert eval_typing(CondIsIntLiteral[str]) == Literal[False] + + +def test_eval_cond_02(): + # Only the selected branch is evaluated + d = eval_typing(_Cond[Literal[True], int, RaiseError[Literal["nope"]]]) + assert d is int + + with pytest.raises(TypeMapError, match="nope"): + eval_typing(_Cond[Literal[False], int, RaiseError[Literal["nope"]]]) + + +def test_eval_bool_ops_01(): + t, f = Literal[True], Literal[False] + nope = RaiseError[Literal["nope"]] + + assert eval_typing(_Not[t]) == f + assert eval_typing(_Not[int]) == t + assert eval_typing(_And[t, IsAssignable[bool, int]]) == t + assert eval_typing(_And[t, f, t]) == f + assert eval_typing(_Or[f, IsAssignable[int, bool]]) == f + assert eval_typing(_Or[f, t]) == t + + # Short circuiting + assert eval_typing(_And[f, nope]) == f + assert eval_typing(_Or[t, nope]) == t + with pytest.raises(TypeMapError, match="nope"): + eval_typing(_And[t, nope]) + + # Substitution + type NotInt[X] = _Not[IsAssignable[X, int]] + assert eval_typing(NotInt[str]) == t + assert eval_typing(NotInt[bool]) == f + + +def test_eval_bool_ops_02(): + t, f = Literal[True], Literal[False] + type AnyInt[Ts] = _Any[_UnpackMap[Ts, T, IsAssignable[T, int]]] + type AllInt[Ts] = _All[_UnpackMap[Ts, T, IsAssignable[T, int]]] + + assert eval_typing(AnyInt[tuple[str, bool]]) == t + assert eval_typing(AnyInt[tuple[str, float]]) == f + assert eval_typing(AnyInt[tuple[()]]) == f + assert eval_typing(AllInt[tuple[int, bool]]) == t + assert eval_typing(AllInt[tuple[int, str]]) == f + assert eval_typing(AllInt[tuple[()]]) == t + + with pytest.raises(TypeError, match="_UnpackMap"): + _Any[tuple[t]] + with pytest.raises(TypeError, match="_UnpackMap"): + _All[t] + + +type CondLast[T] = _Cond[ + IsEquivalent[Length[T], Literal[1]], + GetArg[T, tuple, Literal[0]], + CondLast[Slice[T, Literal[1], None]], +] + + +def test_eval_cond_03(): + assert eval_typing(CondLast[tuple[int]]) is int + assert eval_typing(CondLast[tuple[int, str, float]]) is float + + +type ListAll[Ts] = tuple[_UnpackMap[Ts, T, list[T]]] +type ListInts[Ts] = tuple[_UnpackMap[Ts, T, list[T], IsAssignable[T, int]]] +type SetAttrs[C] = tuple[ + _UnpackMap[ + Attrs[C], + T, + Member[ + _GetAssociated[T, Literal["name"]], + set[_GetAssociated[T, Literal["type"]]], + ], + ] +] +type SetAttrsComp[C] = tuple[ + *[Member[m.name, set[m.type]] for m in Iter[Attrs[C]]] +] + + +def test_eval_unpack_map_01(): + assert eval_typing(ListAll[tuple[int, str]]) == tuple[list[int], list[str]] + assert eval_typing(ListAll[tuple[()]]) == tuple[()] + assert ( + eval_typing(tuple[bool, _UnpackMap[tuple[int, str], T, T], float]) + == tuple[bool, int, str, float] + ) + assert eval_typing(_UnpackMap[tuple[int, str], T, T] | None) == ( + int | str | None + ) + + +def test_eval_unpack_map_02(): + assert ( + eval_typing(ListInts[tuple[int, str, bool, float]]) + == tuple[list[int], list[bool]] + ) + assert eval_typing(ListInts[tuple[str]]) == tuple[()] + + +def test_eval_unpack_map_03(): + class C: + a: int + b: str + + assert eval_typing(SetAttrs[C]) == eval_typing(SetAttrsComp[C]) + + +def test_eval_unpack_map_04(): + # Only selected elements are evaluated + d = eval_typing( + tuple[ + _UnpackMap[ + tuple[int, str], + T, + _Cond[IsAssignable[T, int], T, RaiseError[Literal["nope"]]], + IsAssignable[T, int], + ] + ] + ) + assert d == tuple[int] + + with pytest.raises(TypeError, match="TypeVar"): + eval_typing(tuple[_UnpackMap[tuple[int], int, int]]) + + +def test_eval_unpack_map_05(): + d = eval_typing(tuple[_UnpackMap[tuple[int, str], T, Callable[[T], T]]]) + assert d == tuple[Callable[[int], int], Callable[[str], str]] + + d = eval_typing( + tuple[_UnpackMap[tuple[int, str], T, Annotated[list[T], "x"]]] + ) + assert d == tuple[Annotated[list[int], "x"], Annotated[list[str], "x"]] + + +def test_eval_unpack_map_06(): + # Nested UnpackMaps work like comprehensions with multiple fors + A = TypeVar("A") + B = TypeVar("B") + + d = eval_typing( + tuple[ + _UnpackMap[ + tuple[int, str], + A, + _UnpackMap[tuple[bool, float], B, tuple[A, B]], + ] + ] + ) + assert d == tuple[*[tuple[a, b] for a in (int, str) for b in (bool, float)]] + + d = eval_typing( + tuple[ + _UnpackMap[ + tuple[tuple[int, str], tuple[()], tuple[bool]], + A, + _UnpackMap[A, B, list[B]], + ] + ] + ) + assert d == tuple[list[int], list[str], list[bool]] + + d = eval_typing( + tuple[ + _UnpackMap[ + tuple[int, str], + A, + _UnpackMap[tuple[int, str], B, tuple[A, B], IsEquivalent[A, B]], + ] + ] + ) + assert d == tuple[tuple[int, int], tuple[str, str]] + + +def test_eval_get_associated_01(): + M = Member[Literal["a"], int] + assert eval_typing(_GetAssociated[M, Literal["name"]]) == Literal["a"] + assert eval_typing(_GetAssociated[M, Literal["type"]]) is int + assert eval_typing(_GetAssociated[M, Literal["quals"]]) is Never + + P = Param[Literal["x"], str] + assert ( + eval_typing(_GetAssociated[P, Literal["kind"]]) + == Literal[ParamKind.POSITIONAL_OR_KEYWORD] + ) + + d = eval_typing( + _GetAssociated[M | Member[Literal["b"], str], Literal["type"]] + ) + assert d == int | str + + class C: + x: int + + d = eval_typing( + _GetAssociated[GetMember[C, Literal["x"]], Literal["definer"]] + ) + assert d is C + + with pytest.raises(TypeMapError, match="no associated type"): + eval_typing(_GetAssociated[int, Literal["name"]]) + with pytest.raises(TypeMapError, match="no associated type"): + eval_typing(_GetAssociated[M, Literal["nope"]]) + + def test_eval_bool_literal_01(): d = eval_typing(_BoolLiteral[True]) assert d == _BoolLiteral[True] diff --git a/typemap/type_eval/_eval_operators.py b/typemap/type_eval/_eval_operators.py index 88c6ec0f..717e1c86 100644 --- a/typemap/type_eval/_eval_operators.py +++ b/typemap/type_eval/_eval_operators.py @@ -24,12 +24,14 @@ Attrs, Bool, Capitalize, + _Cond, DropAnnotations, FromUnion, GenericCallable, GetAnnotations, GetArg, GetArgs, + _GetAssociated, GetMember, GetMemberType, GetSpecialAttr, @@ -53,8 +55,15 @@ Concat, Uncapitalize, UpdateClass, + _UnpackMap, Uppercase, + _All, + _And, + _Any, _BoolLiteral, + _Not, + _Or, + _HasAssociatedTypesGenericAlias, ) ################################################################## @@ -492,6 +501,56 @@ def _eval_Bool(tp, *, ctx): return _eval_bool_tp(tp, ctx) +def _eval_test(tp, ctx) -> bool: + return bool(_eval_bool_tp(_unwrap_anno(_eval_types(tp, ctx)), ctx)) + + +@type_eval.register_evaluator(_Not, lazy=True) +def _eval_Not(tp, *, ctx): + return _BoolLiteral[not _eval_test(tp, ctx)] + + +@type_eval.register_evaluator(_And, lazy=True) +def _eval_And(*tps, ctx): + return _BoolLiteral[all(_eval_test(tp, ctx) for tp in tps)] + + +@type_eval.register_evaluator(_Or, lazy=True) +def _eval_Or(*tps, ctx): + return _BoolLiteral[any(_eval_test(tp, ctx) for tp in tps)] + + +# These aren't lazy, so that the _UnpackMap gets expanded into the arguments +@type_eval.register_evaluator(_Any) +def _eval_Any(*tps, ctx): + return _BoolLiteral[any(_eval_test(tp, ctx) for tp in tps)] + + +@type_eval.register_evaluator(_All) +def _eval_All(*tps, ctx): + return _BoolLiteral[all(_eval_test(tp, ctx) for tp in tps)] + + +@type_eval.register_evaluator(_Cond, lazy=True) +def _eval_Cond(cond, tif, telse, *, ctx): + return _eval_types(tif if _eval_test(cond, ctx) else telse, ctx) + + +@type_eval.register_evaluator(_UnpackMap, lazy=True) +def _eval_UnpackMap(tps, tv, tres, tcond=typing.Literal[True], *, ctx): + if not isinstance(tv, typing.TypeVar): + raise TypeError(f"_UnpackMap variable must be a TypeVar, got {tv}") + + results = [] + for tp in _eval_Iter(tps, ctx=ctx): + if _eval_test(_apply_generic.substitute(tcond, {tv: tp}), ctx): + # _eval_args so that an Unpack or an _UnpackMap gets expanded + results.extend( + _eval_args([_apply_generic.substitute(tres, {tv: tp})], ctx) + ) + return typing.Unpack[tuple[*results]] + + ################################################################## @@ -1146,6 +1205,20 @@ def _eval_GetArgs(tp, base, *, ctx) -> typing.Any: return tuple[*args] # type: ignore[valid-type] +@type_eval.register_evaluator(_GetAssociated) +@_lift_over_unions +def _eval_GetAssociated(tp, name, *, ctx) -> typing.Any: + name = _from_literal(name) + if not ( + isinstance(tp, _HasAssociatedTypesGenericAlias) + and isinstance(getattr(tp.__origin__, name, None), typing.TypeAliasType) + ): + raise TypeMapError( + f"_GetAssociated: {tp!r} has no associated type {name!r}" + ) + return _eval_types(getattr(tp, name), ctx) + + @type_eval.register_evaluator(GetSpecialAttr) @_lift_over_unions def _eval_GetSpecialAttr(tp, attr, *, ctx) -> typing.Any: diff --git a/typemap/type_eval/_eval_typing.py b/typemap/type_eval/_eval_typing.py index 9830bf68..059a9816 100644 --- a/typemap/type_eval/_eval_typing.py +++ b/typemap/type_eval/_eval_typing.py @@ -34,14 +34,18 @@ class StuckException(Exception): _eval_funcs: dict[type, typing.Callable[..., Any]] = {} +# Evaluators that receive their arguments unevaluated +_lazy_eval_funcs: set[type] = set() def register_evaluator[T: typing.Callable[..., Any]]( - typ: type, + typ: type, *, lazy: bool = False ) -> typing.Callable[[T], T]: def func(f: T) -> T: assert typ not in _eval_funcs _eval_funcs[typ] = f + if lazy: + _lazy_eval_funcs.add(typ) return f return func @@ -427,6 +431,9 @@ def _eval_applied_class(obj: typing_GenericAlias, ctx: EvalContext): """Eval a typing._GenericAlias -- an applied user-defined class""" # generic *classes* are typing._GenericAlias while generic type # aliases are types.GenericAlias? Why in the world. + if obj.__origin__ in _lazy_eval_funcs: + return _eval_funcs[obj.__origin__](*typing.get_args(obj), ctx=ctx) + new_args = _eval_args(typing.get_args(obj), ctx) if func := _eval_funcs.get(obj.__origin__): @@ -461,6 +468,4 @@ def _eval_ty_or_list(obj): @_eval_types_impl.register def _eval_union(obj: typing.Union, ctx: EvalContext): - args: typing.Sequence[typing.Any] = obj.__args__ - new_args = tuple(_eval_types(arg, ctx) for arg in args) - return typing.Union[new_args] + return typing.Union[_eval_args(obj.__args__, ctx)] diff --git a/typemap/typing.py b/typemap/typing.py index ba106e58..e1c05618 100644 --- a/typemap/typing.py +++ b/typemap/typing.py @@ -290,6 +290,16 @@ class GetSpecialAttr[T: type, Attr: str]: pass +class _GetAssociated[T, Name: str]: + """_GetAssociated[T, Name]: Fetch the associated type Name of T. + + Equivalent to ``T.``, for classes decorated with + has_associated_types (like Member and Param). + """ + + pass + + class Length[S: tuple]: pass @@ -393,6 +403,56 @@ def Bool(self, tp): return _BoolGenericAlias(self, tp) +@_SpecialForm +def _Not(self, tp): + return _BoolGenericAlias(self, tp) + + +@_SpecialForm +def _And(self, tps): + return _BoolGenericAlias(self, tps) + + +@_SpecialForm +def _Or(self, tps): + return _BoolGenericAlias(self, tps) + + +@_SpecialForm +def _Any(self, tp): + """_Any[_UnpackMap[...]]: Equivalent to ``any(...)``""" + if typing.get_origin(tp) is not _UnpackMap: + raise TypeError(f"_Any expected an _UnpackMap, got {tp!r}") + return _BoolGenericAlias(self, tp) + + +@_SpecialForm +def _All(self, tp): + """_All[_UnpackMap[...]]: Equivalent to ``all(...)``""" + if typing.get_origin(tp) is not _UnpackMap: + raise TypeError(f"_All expected an _UnpackMap, got {tp!r}") + return _BoolGenericAlias(self, tp) + + +class _Cond[B, T, F]: + """_Cond[B, T, F]: Equivalent to ``T if Bool[B] else F``. + + Only the selected branch is evaluated. + """ + + pass + + +class _UnpackMap[Ts, T, R, C = Literal[True]]: + """_UnpackMap[Ts, T, R, C]: Equivalent to ``*[R for T in Iter[Ts] if C]``. + + T must be a TypeVar, which is bound to each element of Ts in turn + while evaluating R and C. + """ + + pass + + class _BoolLiteralGenericAlias(_LiteralGenericAlias, _root=True): def __bool__(self): return typing.get_args(self)[0]