Add _Cond, _UnpackMap, _GetAssociated to represent our syntax

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]

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:

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)]