Support star unpacking of operators, aliases, and tuples

pep.rst

@@ -1093,11 +1093,6 @@ NumPy-style broadcasting
         ) -> Array[DType, *Broadcast[tuple[*Shape], tuple[*Shape2]]]:
             raise BaseException
 
-    type AppendTuple[A, B] = tuple[
-        *[x for x in typing.Iter[A]],
-        B,
-    ]
-
     type MergeOne[T, S] = (
         T
         if typing.Matches[T, S] or typing.Matches[S, Literal[1]]
@@ -1118,8 +1113,9 @@ NumPy-style broadcasting
         if typing.Bool[Empty[T]]
         else T
         if typing.Bool[Empty[S]]
-        else AppendTuple[
-            Broadcast[DropLast[T], DropLast[S]], MergeOne[Last[T], Last[S]]
+        else tuple[
+            *Broadcast[DropLast[T], DropLast[S]],
+            MergeOne[Last[T], Last[S]],
         ]
     )
 

tests/test_fastapilike_1.py

@@ -67,40 +67,12 @@ class _Default:
     ],
     Literal["ClassVar"],
 ]
-type AddInit[T] = NewProtocol[
-    InitFnType[T],
-    *[x for x in Iter[Members[T]]],
-]
-
-"""TODO:
-
-We would really like to instead write:
 
 type AddInit[T] = NewProtocol[
     InitFnType[T],
     *Members[T],
 ]
 
-but we struggle here because typing wants to unpack the Members tuple
-itself.  I'm not sure if there is a nice way to resolve this. We
-*could* make our consumers (NewProtocol etc) be more flexible about
-these things but I don't think that is right.
-
-The frustrating thing is that it doesn't do much with the unpacked
-version, just some checks!
-
-We could fix typing to allow it, and probably provide a hack around it
-in the mean time.
-
-Lurr! Writing *this* gets past the typing checks (though we don't
-support it yet):
-
-type AddInit[T] = NewProtocol[
-    InitFnType[T],
-    *tuple[*Members[T]],
-]
-"""
-
 # Strip `| None` from a type by iterating over its union components
 # and filtering
 type NotOptional[T] = Union[

tests/test_nplike.py

@@ -14,11 +14,6 @@ def __add__[*Shape2](
         raise BaseException
 
 
-type AppendTuple[A, B] = tuple[
-    *[x for x in typing.Iter[A]],
-    B,
-]
-
 type MergeOne[T, S] = (
     T
     if typing.Matches[T, S] or typing.Matches[S, Literal[1]]
@@ -39,8 +34,9 @@ def __add__[*Shape2](
     if typing.Bool[Empty[T]]
     else T
     if typing.Bool[Empty[S]]
-    else AppendTuple[
-        Broadcast[DropLast[T], DropLast[S]], MergeOne[Last[T], Last[S]]
+    else tuple[
+        *Broadcast[DropLast[T], DropLast[S]],
+        MergeOne[Last[T], Last[S]],
     ]
 )
 
@@ -49,11 +45,7 @@ def __add__[*Shape2](
 type GetElem[T] = typing.GetArg[T, Array, Literal[0]]
 type GetShape[T] = typing.Slice[typing.GetArgs[T, Array], Literal[1], None]
 
-# type Apply[T, S] = Array[GetElem[T], *Broadcast[GetShape[T], GetShape[S]]]
-type Apply[T, S] = Array[
-    GetElem[T],
-    *[x for x in typing.Iter[Broadcast[GetShape[T], GetShape[S]]]],
-]
+type Apply[T, S] = Array[GetElem[T], *Broadcast[GetShape[T], GetShape[S]]]
 
 ######
 from typemap.type_eval import eval_typing, TypeMapError

tests/test_type_eval.py

@@ -74,10 +74,7 @@ class F_int(F[int]):
             if not IsSub[GetType[p], A]
             else Member[GetName[p], OrGotcha[MapRecursive[A]]]
         )
-        # XXX: This next line *ought* to work, but we haven't
-        # implemented it yet.
-        # for p in Iter[*Attrs[A], *Attrs[F_int]]
-        for p in Iter[ConcatTuples[Attrs[A], Attrs[F_int]]]
+        for p in Iter[tuple[*Attrs[A], *Attrs[F_int]]]
     ],
     Member[Literal["control"], float],
 ]

typemap/type_eval/_eval_operators.py

@@ -1088,7 +1088,8 @@ def _eval_NewProtocol(*etyps: Member, ctx):
     dct: dict[str, object] = {}
     dct["__annotations__"] = annos = {}
 
-    for tname, typ, quals, init, _ in (typing.get_args(prop) for prop in etyps):
+    members = [typing.get_args(prop) for prop in etyps]
+    for tname, typ, quals, init, _ in members:
         name = _eval_literal(tname, ctx)
         typ = _eval_types(typ, ctx)
         tquals = _eval_types(quals, ctx)

typemap/type_eval/_eval_typing.py

@@ -15,11 +15,12 @@
     _CallableGenericAlias as typing_CallableGenericAlias,
     _LiteralGenericAlias as typing_LiteralGenericAlias,
     _AnnotatedAlias as typing_AnnotatedAlias,
+    _UnpackGenericAlias as typing_UnpackGenericAlias,
 )
 
 
 if typing.TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Sequence
 
 from . import _apply_generic, _typing_inspect
 
@@ -377,14 +378,38 @@ def _eval_type_alias(obj: typing.TypeAliasType, ctx: EvalContext):
     return _eval_types(unpacked, ctx)
 
 
+def _eval_args(args: Sequence[Any], ctx: EvalContext) -> tuple[Any]:
+    evaled = []
+    for arg in args:
+        ev = _eval_types(arg, ctx)
+        if isinstance(ev, typing_UnpackGenericAlias):
+            if (args := ev.__typing_unpacked_tuple_args__) is not None:
+                evaled.extend(args)
+            else:
+                evaled.append(ev)
+        else:
+            evaled.append(ev)
+    return tuple(evaled)
+
+
 @_eval_types_impl.register
 def _eval_applied_type_alias(obj: types.GenericAlias, ctx: EvalContext):
     """Eval a types.GenericAlias -- typically an applied type alias
 
     This is typically an application of a type alias... except it can
-    also be an application of a built-in type (like list, tuple, dict)
+    also be an application of a built-in type (like list, tuple, dict).
+
+    It can *also* have an Unpack integrated with it, if __unpacked__ is set.
     """
-    new_args = tuple(_eval_types(arg, ctx) for arg in obj.__args__)
+
+    # If __unpacked__ is set, then we reconstruct a version without
+    # __unpacked__ set and evaluate *that*. This centralizes the
+    # unpacked handling and simplifies the cache situation.
+    if obj.__unpacked__:
+        stripped = _apply_type(obj.__origin__, obj.__args__)
+        return typing.Unpack[_eval_types(stripped, ctx)]
+
+    new_args = _eval_args(obj.__args__, ctx)
 
     new_obj = _apply_type(obj.__origin__, new_args)
     if isinstance(obj.__origin__, type):
@@ -433,7 +458,7 @@ 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.
-    new_args = tuple(_eval_types(arg, ctx) for arg in typing.get_args(obj))
+    new_args = _eval_args(typing.get_args(obj), ctx)
 
     if func := _eval_funcs.get(obj.__origin__):
         ret = func(*new_args, ctx=ctx)

typemap/typing.py

@@ -1,10 +1,56 @@
 import contextvars
 import typing
-from typing import Literal
-from typing import _GenericAlias, _LiteralGenericAlias  # type: ignore
+import types
+
+from typing import Literal, Unpack
+from typing import _GenericAlias, _LiteralGenericAlias, _UnpackGenericAlias  # type: ignore
 
 _SpecialForm: typing.Any = typing._SpecialForm
 
+###
+
+# Here is a bunch of annoying internals stuff!
+
+
+class _TupleLikeOperator:
+    @classmethod
+    def __class_getitem__(cls, args):
+        # Return an _IterSafeGenericAlias instead of a _GenericAlias
+        res = super().__class_getitem__(args)
+        return _IterSafeGenericAlias(res.__origin__, res.__args__)
+
+
+# The base _GenericAlias has an __iter__ method that returns
+# Unpack[self], which blows up when it's passed to something and
+# doesn't have a tuple inside (because it hasn't been evaluated yet!).
+# So we make own _GenericAlias that makes our own _UnpackGenericAlias
+# that we make sure works.
+#
+# Probably these exact hacks will need to go into our
+# typing_extensions version of this, but for the typing version they
+# can get merged into real classes.
+class _IterSafeGenericAlias(_GenericAlias, _root=True):  # type: ignore[call-arg]
+    def __iter__(self):
+        yield _IterSafeUnpackGenericAlias(origin=Unpack, args=(self,))
+
+
+class _IterSafeUnpackGenericAlias(_UnpackGenericAlias, _root=True):  # type: ignore[call-arg]
+    @property
+    def __typing_unpacked_tuple_args__(self):
+        # This is basically the same as in _UnpackGenericAlias except
+        # we don't blow up if the origin isn't a tuple.
+        assert self.__origin__ is Unpack
+        assert len(self.__args__) == 1
+        (arg,) = self.__args__
+        if isinstance(arg, (_GenericAlias, types.GenericAlias)):
+            if arg.__origin__ is tuple:
+                return arg.__args__
+        return None
+
+
+###
+
+
 # Not type-level computation but related
 
 
@@ -119,15 +165,15 @@ class Param[N: str | None, T, Q: ParamQuals = typing.Never]:
 type GetDefiner[T: Member] = GetMemberType[T, Literal["definer"]]
 
 
-class Attrs[T]:
+class Attrs[T](_TupleLikeOperator):
     pass
 
 
-class Members[T]:
+class Members[T](_TupleLikeOperator):
     pass
 
 
-class FromUnion[T]:
+class FromUnion[T](_TupleLikeOperator):
     pass
 
 
@@ -143,7 +189,7 @@ class GetArg[Tp, Base, Idx: int]:
     pass
 
 
-class GetArgs[Tp, Base]:
+class GetArgs[Tp, Base](_TupleLikeOperator):
     pass
 
 
@@ -155,7 +201,9 @@ class Length[S: tuple]:
     pass
 
 
-class Slice[S: str | tuple, Start: int | None, End: int | None]:
+class Slice[S: str | tuple, Start: int | None, End: int | None](
+    _TupleLikeOperator
+):
     pass