PEP 827: Swap out the numpy-like example for zip

Author: msullivan

peps/pep-0827.rst

@@ -1079,8 +1079,8 @@ based on iterating over all attributes.
     type InitFnType[T] = typing.Member[
         Literal["__init__"],
         Callable[
-            [
-                typing.Param[Literal["self"], Self],
+            typing.Params[
+                typing.Param[Literal["self"], T],
                 *[
                     typing.Param[
                         p.name,
@@ -1117,7 +1117,7 @@ based on iterating over all attributes.
         # Add the computed __init__ function
         InitFnType[T],
     ]:
-        pass
+        raise NotImplementedError
 
 Or to create a base class (a la Pydantic) that does.
 
@@ -1130,86 +1130,40 @@ Or to create a base class (a la Pydantic) that does.
             # Add the computed __init__ function
             InitFnType[T],
         ]:
-            super().__init_subclass__()
-
-
-NumPy-style broadcasting
-------------------------
-
-One of the motivations for the introduction of ``TypeVarTuple`` in
-:pep:`646` is to represent the shapes of multi-dimensional
-arrays, such as::
-
-  x: Array[float, L[480], L[640]] = Array()
-
-The example in that PEP shows how ``TypeVarTuple`` can be used to
-make sure that both sides of an arithmetic operation have matching
-shapes. Most multi-dimensional array libraries, however, also support
-`broadcasting <#broadcasting_>`__, which allows the mixing of differently
-shaped data.  With this PEP, we can define a ``Broadcast[A, B]`` type
-alias, and then use it as a return type::
-
-    class Array[DType, *Shape]:
-        def __add__[*Shape2](
-            self,
-            other: Array[DType, *Shape2]
-        ) -> Array[DType, *Broadcast[tuple[*Shape], tuple[*Shape2]]]:
-            raise BaseException
-
-(The somewhat clunky syntax of wrapping the ``TypeVarTuple`` in
-another ``tuple`` is because typecheckers currently disallow having
-two ``TypeVarTuple`` arguments. A possible improvement would be to
-allow writing the bare (non-starred or ``Unpack``-ed) variable name to
-mean its interpretation as a tuple.)
+            pass
 
-We can then do::
 
-    a1: Array[float, L[4], L[1]]
-    a2: Array[float, L[3]]
-    a1 + a2  # Array[builtins.float, Literal[4], Literal[3]]
+.. _pep827-zip-impl:
 
-    b1: Array[float, int, int]
-    b2: Array[float, int]
-    b1 + b2  # Array[builtins.float, int, int]
-
-    err1: Array[float, L[4], L[2]]
-    err2: Array[float, L[3]]
-    # err1 + err2  # E: Broadcast mismatch: Literal[2], Literal[3]
-
-
-Note that this is meant to be an example of the expressiveness of type
-manipulation, and not any kind of final proposal about the typing of
-tensor types.
+zip-like functions
+------------------
 
-.. _pep827-numpy-impl:
-
-Implementation
-''''''''''''''
+Using type iteration and ``GetArg``, we can give a proper type to ``zip``.
 
 ::
 
-    class Array[DType, *Shape]:
-        def __add__[*Shape2](
-            self, other: Array[DType, *Shape2]
-        ) -> Array[DType, *Broadcast[tuple[*Shape], tuple[*Shape2]]]:
-            raise BaseException
+    type ElemOf[T] = typing.GetArg[T, Iterable, Literal[0]]
+
+    def zip[*Ts](
+        *args: *Ts, strict: bool = False
+    ) -> Iterator[tuple[*[ElemOf[t] for t in typing.Iter[tuple[*Ts]]]]]:
+        return builtins.zip(*args, strict=strict)  # type: ignore[call-overload]
 
-``MergeOne`` is the core of the broadcasting operation. If the two types
-are equivalent, we take the first, and if either of the types is
-``Literal[1]`` then we take the other.
+Using the ``Slice`` operator and type alias recursion, we can
+also give a more precise type for zipping together heterogeneous tuples.
 
-On a mismatch, we use the ``RaiseError`` operator to produce an error
-message identifying the two types.
+For example, zipping ``tuple[int, str]`` and ``tuple[str, bool]``
+should produce ``tuple[tuple[int, float], tuple[str, bool]]``
 
 ::
 
-    type MergeOne[T, S] = (
-        T
-        if typing.IsEquivalent[T, S] or typing.IsEquivalent[S, Literal[1]]
-        else S
-        if typing.IsEquivalent[T, Literal[1]]
-        else typing.RaiseError[Literal["Broadcast mismatch"], T, S]
-    )
+    def zip_pairs[*Ts, *Us](
+        a: tuple[*Ts], b: tuple[*Us]
+    ) -> Zip[tuple[*Ts], tuple[*Us]]:
+        return cast(
+            Zip[tuple[*Ts], tuple[*Us]],
+            tuple(zip(a, b, strict=True)),
+        )
 
     type DropLast[T] = typing.Slice[T, Literal[0], Literal[-1]]
     type Last[T] = typing.GetArg[T, tuple, Literal[-1]]
@@ -1218,20 +1172,18 @@ message identifying the two types.
     # recursions when T is not a tuple.
     type Empty[T] = typing.IsAssignable[typing.Length[T], Literal[0]]
 
-Broadcast recursively walks down the input tuples applying ``MergeOne``
-until one of them is empty.
+Zip recursively walks down the input tuples until one or both of them
+is empty. If the lengths don't match (because only one is empty),
+raise an error.
 
 ::
 
-    type Broadcast[T, S] = (
-        S
-        if typing.Bool[Empty[T]]
-        else T
-        if typing.Bool[Empty[S]]
-        else tuple[
-            *Broadcast[DropLast[T], DropLast[S]],
-            MergeOne[Last[T], Last[S]],
-        ]
+    type Zip[T, S] = (
+        tuple[()]
+        if typing.Bool[Empty[T]] and typing.Bool[Empty[S]]
+        else typing.RaiseError[Literal["Zip length mismatch"], T, S]
+        if typing.Bool[Empty[T]] or typing.Bool[Empty[S]]
+        else tuple[*Zip[DropLast[T], DropLast[S]], tuple[Last[T], Last[S]]]
     )