Repository: maxfischer2781/asyncstdlib

Issue:
Add zip strict flag
Python 3.10 / [PEP 618](https://www.python.org/dev/peps/pep-0618/) will add a ``strict`` flag to ``zip``. When ``strict=True``, an exception is raised if the iterables are not of the same length.

Comments:

Patch:
diff --git a/asyncstdlib/builtins.py b/asyncstdlib/builtins.py
index 4091b8c..06eb44e 100644
--- a/asyncstdlib/builtins.py
+++ b/asyncstdlib/builtins.py
@@ -134,10 +134,12 @@ async def any(iterable: AnyIterable[T]) -> bool:
         return False
 
 
-async def zip(*iterables: AnyIterable[T]) -> AsyncIterator[Tuple[T, ...]]:
+async def zip(*iterables: AnyIterable[T], strict=False) -> AsyncIterator[Tuple[T, ...]]:
     """
     Create an async iterator that aggregates elements from each of the (async) iterables
 
+    :raises ValueError: if the ``iterables`` are not equal length and ``strict`` is set
+
     The next element of ``zip`` is a :py:class:`tuple` of the next element of
     each of its ``iterables``. As soon as any of its ``iterables`` is exhausted,
     ``zip`` is exhausted as well. This means that if ``zip`` receives *n* ``iterables``,
@@ -151,16 +153,19 @@ async def zip(*iterables: AnyIterable[T]) -> AsyncIterator[Tuple[T, ...]]:
 
     If ``iterables`` is empty, the ``zip`` iterator is empty as well.
     Multiple ``iterables`` may be mixed regular and async iterables.
+
+    When called with ``strict=True``, all ``iterables`` must be of same length;
+    in this mode ``zip`` raises :py:exc:`ValueError` if any ``iterables`` are not
+    exhausted with the others.
     """
     if not iterables:
         return
     aiters = (*(aiter(it) for it in iterables),)
     del iterables
     try:
-        while True:
-            yield (*[await anext(it) for it in aiters],)
-    except StopAsyncIteration:
-        return
+        inner = _zip_inner(aiters) if not strict else _zip_inner_strict(aiters)
+        async for items in inner:
+            yield items
     finally:
         for iterator in aiters:
             try:
@@ -171,6 +176,40 @@ async def zip(*iterables: AnyIterable[T]) -> AsyncIterator[Tuple[T, ...]]:
                 await aclose
 
 
+async def _zip_inner(aiters):
+    try:
+        while True:
+            yield (*[await anext(it) for it in aiters],)
+    except StopAsyncIteration:
+        return
+
+
+async def _zip_inner_strict(aiters):
+    tried = 0
+    try:
+        while True:
+            items = []
+            for tried, _aiter in _sync_builtins.enumerate(aiters):  # noqa: B007
+                items.append(await anext(_aiter))
+            yield (*items,)
+    except StopAsyncIteration:
+        # after the first iterable provided an item, some later iterable was empty
+        if tried > 0:
+            plural = " " if tried == 1 else "s 1-"
+            raise ValueError(
+                f"zip() argument {tried+1} is shorter than argument{plural}{tried}"
+            )
+        # after the first iterable was empty, some later iterable may be not
+        sentinel = object()
+        for tried, _aiter in _sync_builtins.enumerate(aiters):
+            if await anext(_aiter, sentinel) is not sentinel:
+                plural = " " if tried == 1 else "s 1-"
+                raise ValueError(
+                    f"zip() argument {tried+1} is longer than argument{plural}{tried}"
+                )
+        return
+
+
 class SyncVariadic(Protocol[T, R]):
     """Type of a ``def`` function taking any number of arguments"""
 
diff --git a/docs/source/api/builtins.rst b/docs/source/api/builtins.rst
index 3cfc71b..f4c4c42 100644
--- a/docs/source/api/builtins.rst
+++ b/docs/source/api/builtins.rst
@@ -38,9 +38,13 @@ Iterator transforming
 .. autofunction:: filter(function: (T) → (await) bool, iterable: (async) iter T)
     :async-for: :T
 
-.. autofunction:: zip(*iterables: (async) iter T)
+.. autofunction:: zip(*iterables: (async) iter T, strict: bool = True)
     :async-for: :(T, ...)
 
+    .. versionadded:: 3.10.0
+
+        The ``strict`` parameter.
+
 .. autofunction:: map(function: (T, ...) → (await) R, iterable: (async) iter T, ...)
     :async-for: :R
 


Test patch:
diff --git a/unittests/test_builtins.py b/unittests/test_builtins.py
index 776867e..18ce9c7 100644
--- a/unittests/test_builtins.py
+++ b/unittests/test_builtins.py
@@ -57,6 +57,21 @@ async def test_zip():
         assert False
 
 
+@sync
+async def test_zip_strict():
+    async for va, vb in a.zip(asyncify(range(5)), range(5), strict=True):
+        assert va == vb
+    with pytest.raises(ValueError):
+        async for _ in a.zip(asyncify(range(5)), range(6), strict=True):
+            pass
+    with pytest.raises(ValueError):
+        async for _ in a.zip(asyncify(range(6)), range(5), strict=True):
+            pass
+    with pytest.raises(ValueError):
+        async for _ in a.zip(*[range(5)] * 6, range(6), strict=True):
+            pass
+
+
 @sync
 async def test_zip_close_immediately():
     closed = False