From a36a6e5f7704b520b0df3ab38253e74225cdb2d5 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 29 Apr 2020 15:15:57 -0700
Subject: [PATCH 01/26] Implement zip(..., total=True)

---
 Lib/test/test_builtin.py   |  2 +-
 Lib/test/test_itertools.py |  8 ++---
 Python/bltinmodule.c       | 66 +++++++++++++++++++++++++++++++-------
 3 files changed, 59 insertions(+), 17 deletions(-)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index 290ba2cad8e5eb..644bc06b872781 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1569,7 +1569,7 @@ def test_zip_pickle(self):
         a = (1, 2, 3)
         b = (4, 5, 6)
         t = [(1, 4), (2, 5), (3, 6)]
-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
             z1 = zip(a, b)
             self.check_iter_pickle(z1, t, proto)
 
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index eaa6197bec395c..241b93d0686eb1 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -903,17 +903,17 @@ def test_zip_tuple_reuse(self):
         ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
         self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
 
-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
             ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count()), proto))]
             self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
 
-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
             testIntermediate = zip('abc',count())
             next(testIntermediate)
             ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate, proto))]
             self.assertEqual(ans, [('b', 1), ('c', 2)])
 
-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
             self.pickletest(proto, zip('abc', count()))
 
     def test_ziplongest(self):
@@ -1202,7 +1202,7 @@ def test_starmap(self):
         c = starmap(operator.pow, zip(range(3), range(1,7)))
         self.assertEqual(list(copy.deepcopy(c)), ans)
 
-        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
             c = starmap(operator.pow, zip(range(3), range(1,7)))
             self.pickletest(proto, c)
 
diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 14c3e96fb51ab7..ea0d55fcf2765f 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2521,8 +2521,8 @@ builtin_issubclass_impl(PyObject *module, PyObject *cls,
 
 typedef struct {
     PyObject_HEAD
-    Py_ssize_t          tuplesize;
-    PyObject *ittuple;                  /* tuple of iterators */
+    Py_ssize_t tuplesize;  /* negative when total=True */
+    PyObject *ittuple;     /* tuple of iterators */
     PyObject *result;
 } zipobject;
 
@@ -2534,9 +2534,21 @@ zip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     PyObject *ittuple;  /* tuple of iterators */
     PyObject *result;
     Py_ssize_t tuplesize;
+    int total = 0;
 
-    if (type == &PyZip_Type && !_PyArg_NoKeywords("zip", kwds))
-        return NULL;
+    if (kwds) {
+        PyObject *empty = PyTuple_New(0);
+        if (empty == NULL) {
+            return NULL;
+        }
+        static char *kwlist[] = {"total", NULL};
+        int parsed = PyArg_ParseTupleAndKeywords(
+                empty, kwds, "|$p:zip", kwlist, &total);
+        Py_DECREF(empty);
+        if (!parsed) {
+            return NULL;
+        }
+    }
 
     /* args must be a tuple */
     assert(PyTuple_Check(args));
@@ -2575,7 +2587,7 @@ zip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         return NULL;
     }
     lz->ittuple = ittuple;
-    lz->tuplesize = tuplesize;
+    lz->tuplesize = total ? -tuplesize : tuplesize;
     lz->result = result;
 
     return (PyObject *)lz;
@@ -2602,7 +2614,7 @@ static PyObject *
 zip_next(zipobject *lz)
 {
     Py_ssize_t i;
-    Py_ssize_t tuplesize = lz->tuplesize;
+    Py_ssize_t tuplesize = Py_ABS(lz->tuplesize);
     PyObject *result = lz->result;
     PyObject *it;
     PyObject *item;
@@ -2617,6 +2629,9 @@ zip_next(zipobject *lz)
             item = (*Py_TYPE(it)->tp_iternext)(it);
             if (item == NULL) {
                 Py_DECREF(result);
+                if (lz->tuplesize < 0 && !PyErr_Occurred()) {
+                    goto check;
+                }
                 return NULL;
             }
             olditem = PyTuple_GET_ITEM(result, i);
@@ -2632,33 +2647,60 @@ zip_next(zipobject *lz)
             item = (*Py_TYPE(it)->tp_iternext)(it);
             if (item == NULL) {
                 Py_DECREF(result);
+                if (lz->tuplesize < 0 && !PyErr_Occurred()) {
+                    goto check;
+                }
                 return NULL;
             }
             PyTuple_SET_ITEM(result, i, item);
         }
     }
     return result;
+check:
+    if (i) {
+        return PyErr_Format(PyExc_ValueError, "%s() argument %d is too short",
+                            Py_TYPE(lz)->tp_name, i + 1);
+    }
+    for (i = 1; i < tuplesize; i++) {
+        it = PyTuple_GET_ITEM(lz->ittuple, i);
+        item = (*Py_TYPE(it)->tp_iternext)(it);
+        if (item) {
+            Py_DECREF(item);
+            return PyErr_Format(PyExc_ValueError, "%s() argument %d is too long",
+                                Py_TYPE(lz)->tp_name, i + 1);
+        }
+        if (PyErr_Occurred()) {
+            return NULL;
+        }
+    }
+    return NULL;
 }
 
 static PyObject *
-zip_reduce(zipobject *lz, PyObject *Py_UNUSED(ignored))
+zip_getnewargs_ex(zipobject *lz, PyObject *Py_UNUSED(ignored))
 {
     /* Just recreate the zip with the internal iterator tuple */
-    return Py_BuildValue("OO", Py_TYPE(lz), lz->ittuple);
+    if (lz->tuplesize < 0) {
+        return Py_BuildValue("O{sO}", lz->ittuple, "total", Py_True);
+    }
+    return Py_BuildValue("O{}", lz->ittuple);
 }
 
 static PyMethodDef zip_methods[] = {
-    {"__reduce__",   (PyCFunction)zip_reduce,   METH_NOARGS, reduce_doc},
-    {NULL,           NULL}           /* sentinel */
+    {"__getnewargs_ex__", (PyCFunction)zip_getnewargs_ex, METH_NOARGS, NULL},
+    {NULL}  /* sentinel */
 };
 
 PyDoc_STRVAR(zip_doc,
-"zip(*iterables) --> zip object\n\
+"zip(*iterables[, total=False]) --> zip object\n\
 \n\
 Return a zip object whose .__next__() method returns a tuple where\n\
 the i-th element comes from the i-th iterable argument.  The .__next__()\n\
 method continues until the shortest iterable in the argument sequence\n\
-is exhausted and then it raises StopIteration.");
+is exhausted and then it raises StopIteration.\n\
+\n\
+If total is True, raise a ValueError if one of the arguments is exhausted\n\
+before the others.");
 
 PyTypeObject PyZip_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)

From b59fe70aecd0b93474c5e946ad879b0e3622d22b Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 30 Apr 2020 16:58:54 -0700
Subject: [PATCH 02/26] total -> strict

---
 Python/bltinmodule.c | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index ea0d55fcf2765f..f27e42519f556c 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2521,7 +2521,7 @@ builtin_issubclass_impl(PyObject *module, PyObject *cls,
 
 typedef struct {
     PyObject_HEAD
-    Py_ssize_t tuplesize;  /* negative when total=True */
+    Py_ssize_t tuplesize;  /* negative when strict=True */
     PyObject *ittuple;     /* tuple of iterators */
     PyObject *result;
 } zipobject;
@@ -2534,16 +2534,16 @@ zip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     PyObject *ittuple;  /* tuple of iterators */
     PyObject *result;
     Py_ssize_t tuplesize;
-    int total = 0;
+    int strict = 0;
 
     if (kwds) {
         PyObject *empty = PyTuple_New(0);
         if (empty == NULL) {
             return NULL;
         }
-        static char *kwlist[] = {"total", NULL};
+        static char *kwlist[] = {"strict", NULL};
         int parsed = PyArg_ParseTupleAndKeywords(
-                empty, kwds, "|$p:zip", kwlist, &total);
+                empty, kwds, "|$p:zip", kwlist, &strict);
         Py_DECREF(empty);
         if (!parsed) {
             return NULL;
@@ -2587,7 +2587,7 @@ zip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         return NULL;
     }
     lz->ittuple = ittuple;
-    lz->tuplesize = total ? -tuplesize : tuplesize;
+    lz->tuplesize = strict ? -tuplesize : tuplesize;
     lz->result = result;
 
     return (PyObject *)lz;
@@ -2666,7 +2666,8 @@ zip_next(zipobject *lz)
         item = (*Py_TYPE(it)->tp_iternext)(it);
         if (item) {
             Py_DECREF(item);
-            return PyErr_Format(PyExc_ValueError, "%s() argument %d is too long",
+            return PyErr_Format(PyExc_ValueError,
+                                "%s() argument %d is too long",
                                 Py_TYPE(lz)->tp_name, i + 1);
         }
         if (PyErr_Occurred()) {
@@ -2681,7 +2682,7 @@ zip_getnewargs_ex(zipobject *lz, PyObject *Py_UNUSED(ignored))
 {
     /* Just recreate the zip with the internal iterator tuple */
     if (lz->tuplesize < 0) {
-        return Py_BuildValue("O{sO}", lz->ittuple, "total", Py_True);
+        return Py_BuildValue("O{sO}", lz->ittuple, "strict", Py_True);
     }
     return Py_BuildValue("O{}", lz->ittuple);
 }
@@ -2692,15 +2693,15 @@ static PyMethodDef zip_methods[] = {
 };
 
 PyDoc_STRVAR(zip_doc,
-"zip(*iterables[, total=False]) --> zip object\n\
+"zip(*iterables[, strict=False]) --> zip object\n\
 \n\
 Return a zip object whose .__next__() method returns a tuple where\n\
 the i-th element comes from the i-th iterable argument.  The .__next__()\n\
 method continues until the shortest iterable in the argument sequence\n\
 is exhausted and then it raises StopIteration.\n\
 \n\
-If total is True, raise a ValueError if one of the arguments is exhausted\n\
-before the others.");
+If strict is true and one of the arguments is exhausted before the others,\n\
+raise a ValueError.");
 
 PyTypeObject PyZip_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)

From 3ed7f0ab1a6da2f788e1273053964fadd280b532 Mon Sep 17 00:00:00 2001
From: Ram Rachum <ram@rachum.com>
Date: Tue, 21 Apr 2020 18:38:33 +0300
Subject: [PATCH 03/26] Tests for zip(total=True)

---
 Lib/test/test_builtin.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index 644bc06b872781..260d26803298e7 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1585,6 +1585,26 @@ def __iter__(self):
 
         self.assertIs(cm.exception, exception)
 
+    def test_zip_total(self):
+        assert tuple(zip((1, 2, 3), 'abc', total=True)) == (
+            (1, 'a'), (2, 'b'), (3, 'c')
+        )
+        self.assertRaises(ValueError, tuple,
+                          zip((1, 2, 3, 4), 'abc', total=True))
+        self.assertRaises(ValueError, tuple,
+                          zip((1, 2), 'abc', total=True))
+        self.assertRaises(ValueError, tuple,
+                          zip((1, 2), (1, 2), 'abc', total=True))
+
+    def test_zip_total_iterators(self):
+        x = iter(range(5))
+        y = [0]
+        z = iter(range(5))
+        self.assertRaises(ValueError, list,
+                          (zip(x, y, z, total=True)))
+        assert next(x) == 2
+        assert next(z) == 1
+
     def test_format(self):
         # Test the basic machinery of the format() builtin.  Don't test
         #  the specifics of the various formatters

From 3955ca75599ee66c0eb170d77e338fb32f7aea0f Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 30 Apr 2020 22:27:27 -0700
Subject: [PATCH 04/26] assert -> self.assertEqual, total -> strict

---
 Lib/test/test_builtin.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index 260d26803298e7..d61c10d1a298ee 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1585,25 +1585,24 @@ def __iter__(self):
 
         self.assertIs(cm.exception, exception)
 
-    def test_zip_total(self):
-        assert tuple(zip((1, 2, 3), 'abc', total=True)) == (
-            (1, 'a'), (2, 'b'), (3, 'c')
-        )
+    def test_zip_strict(self):
+        self.assertEqual(tuple(zip((1, 2, 3), 'abc', strict=True)),
+                         ((1, 'a'), (2, 'b'), (3, 'c')))
         self.assertRaises(ValueError, tuple,
-                          zip((1, 2, 3, 4), 'abc', total=True))
+                          zip((1, 2, 3, 4), 'abc', strict=True))
         self.assertRaises(ValueError, tuple,
-                          zip((1, 2), 'abc', total=True))
+                          zip((1, 2), 'abc', strict=True))
         self.assertRaises(ValueError, tuple,
-                          zip((1, 2), (1, 2), 'abc', total=True))
+                          zip((1, 2), (1, 2), 'abc', strict=True))
 
-    def test_zip_total_iterators(self):
+    def test_zip_strict_iterators(self):
         x = iter(range(5))
         y = [0]
         z = iter(range(5))
         self.assertRaises(ValueError, list,
-                          (zip(x, y, z, total=True)))
-        assert next(x) == 2
-        assert next(z) == 1
+                          (zip(x, y, z, strict=True)))
+        self.assertEqual(next(x), 2)
+        self.assertEqual(next(z), 1)
 
     def test_format(self):
         # Test the basic machinery of the format() builtin.  Don't test

From e4dfe05260a14b60a66244afec1312a54713bbaa Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Tue, 12 May 2020 11:22:43 -0700
Subject: [PATCH 05/26] Support all pickle protocols.

---
 Lib/test/test_builtin.py   | 2 +-
 Lib/test/test_itertools.py | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index d61c10d1a298ee..a6d2388ab03345 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1569,7 +1569,7 @@ def test_zip_pickle(self):
         a = (1, 2, 3)
         b = (4, 5, 6)
         t = [(1, 4), (2, 5), (3, 6)]
-        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             z1 = zip(a, b)
             self.check_iter_pickle(z1, t, proto)
 
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 241b93d0686eb1..eaa6197bec395c 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -903,17 +903,17 @@ def test_zip_tuple_reuse(self):
         ans = [(x,y) for x, y in copy.deepcopy(zip('abc',count()))]
         self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
 
-        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             ans = [(x,y) for x, y in pickle.loads(pickle.dumps(zip('abc',count()), proto))]
             self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
 
-        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             testIntermediate = zip('abc',count())
             next(testIntermediate)
             ans = [(x,y) for x, y in pickle.loads(pickle.dumps(testIntermediate, proto))]
             self.assertEqual(ans, [('b', 1), ('c', 2)])
 
-        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             self.pickletest(proto, zip('abc', count()))
 
     def test_ziplongest(self):
@@ -1202,7 +1202,7 @@ def test_starmap(self):
         c = starmap(operator.pow, zip(range(3), range(1,7)))
         self.assertEqual(list(copy.deepcopy(c)), ans)
 
-        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             c = starmap(operator.pow, zip(range(3), range(1,7)))
             self.pickletest(proto, c)
 

From 4d41e21a53404f5c34b580ce61c0f5871570737f Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Tue, 12 May 2020 11:24:11 -0700
Subject: [PATCH 06/26] Use __reduce__/__setstate__ for pickling.

---
 Python/bltinmodule.c | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index f27e42519f556c..c77f7c2a931c90 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2678,17 +2678,31 @@ zip_next(zipobject *lz)
 }
 
 static PyObject *
-zip_getnewargs_ex(zipobject *lz, PyObject *Py_UNUSED(ignored))
+zip_reduce(zipobject *lz, PyObject *Py_UNUSED(ignored))
 {
     /* Just recreate the zip with the internal iterator tuple */
     if (lz->tuplesize < 0) {
-        return Py_BuildValue("O{sO}", lz->ittuple, "strict", Py_True);
+        return Py_BuildValue("OOO", Py_TYPE(lz), lz->ittuple, Py_True);
     }
-    return Py_BuildValue("O{}", lz->ittuple);
+    return Py_BuildValue("OO", Py_TYPE(lz), lz->ittuple);
+}
+
+PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
+
+static PyObject *
+zip_setstate(zipobject *lz, PyObject *state)
+{
+    int strict = PyObject_IsTrue(state);
+    if (strict < 0) {
+        return NULL;
+    }
+    lz->tuplesize = strict ? -Py_ABS(lz->tuplesize) : Py_ABS(lz->tuplesize);
+    Py_RETURN_NONE;
 }
 
 static PyMethodDef zip_methods[] = {
-    {"__getnewargs_ex__", (PyCFunction)zip_getnewargs_ex, METH_NOARGS, NULL},
+    {"__reduce__",   (PyCFunction)zip_reduce,   METH_NOARGS, reduce_doc},
+    {"__setstate__", (PyCFunction)zip_setstate, METH_O,      setstate_doc},
     {NULL}  /* sentinel */
 };
 

From 892961c74c154b6e628934521f5d1fcde7e259f4 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 13 May 2020 08:19:24 -0700
Subject: [PATCH 07/26] Clean up.

---
 Python/bltinmodule.c | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index c77f7c2a931c90..357e571a5714a4 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2666,8 +2666,7 @@ zip_next(zipobject *lz)
         item = (*Py_TYPE(it)->tp_iternext)(it);
         if (item) {
             Py_DECREF(item);
-            return PyErr_Format(PyExc_ValueError,
-                                "%s() argument %d is too long",
+            return PyErr_Format(PyExc_ValueError, "%s() argument %d is too long",
                                 Py_TYPE(lz)->tp_name, i + 1);
         }
         if (PyErr_Occurred()) {

From 2d19de53bfcb95e1dad9bb2c41e91a17b50bb7a5 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Tue, 19 May 2020 18:56:24 -0700
Subject: [PATCH 08/26] Handle StopIteration raised from Pythonland

---
 Python/bltinmodule.c | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 93d28885d164ea..e69739cccb24bc 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2625,7 +2625,7 @@ zip_next(zipobject *lz)
             item = (*Py_TYPE(it)->tp_iternext)(it);
             if (item == NULL) {
                 Py_DECREF(result);
-                if (lz->tuplesize < 0 && !PyErr_Occurred()) {
+                if (lz->tuplesize < 0) {
                     goto check;
                 }
                 return NULL;
@@ -2643,7 +2643,7 @@ zip_next(zipobject *lz)
             item = (*Py_TYPE(it)->tp_iternext)(it);
             if (item == NULL) {
                 Py_DECREF(result);
-                if (lz->tuplesize < 0 && !PyErr_Occurred()) {
+                if (lz->tuplesize < 0) {
                     goto check;
                 }
                 return NULL;
@@ -2652,7 +2652,14 @@ zip_next(zipobject *lz)
         }
     }
     return result;
-check:
+check:;
+    PyObject *err = PyErr_Occurred();
+    if (err) {
+        if (!PyErr_GivenExceptionMatches(err, PyExc_StopIteration)) {
+            return NULL;
+        }
+        PyErr_Clear();
+    }
     if (i) {
         return PyErr_Format(PyExc_ValueError, "%s() argument %d is too short",
                             Py_TYPE(lz)->tp_name, i + 1);

From f2235f24ff285a50238b51a44937cefa8f93c72f Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 20 May 2020 08:43:14 -0700
Subject: [PATCH 09/26] Clean up error handling

---
 Python/bltinmodule.c | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index e69739cccb24bc..9c107713f948da 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2652,15 +2652,12 @@ zip_next(zipobject *lz)
         }
     }
     return result;
-check:;
-    PyObject *err = PyErr_Occurred();
-    if (err) {
-        if (!PyErr_GivenExceptionMatches(err, PyExc_StopIteration)) {
-            return NULL;
-        }
-        PyErr_Clear();
+check:
+    if (PyErr_Occurred() && !PyErr_ExceptionMatches(PyExc_StopIteration)) {
+        return NULL;
     }
     if (i) {
+        PyErr_Clear();
         return PyErr_Format(PyExc_ValueError, "%s() argument %d is too short",
                             Py_TYPE(lz)->tp_name, i + 1);
     }
@@ -2669,6 +2666,7 @@ check:;
         item = (*Py_TYPE(it)->tp_iternext)(it);
         if (item) {
             Py_DECREF(item);
+            PyErr_Clear();
             return PyErr_Format(PyExc_ValueError, "%s() argument %d is too long",
                                 Py_TYPE(lz)->tp_name, i + 1);
         }

From f38ed4f4a427cda379802903a54e2965e426c0b9 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 10:03:40 -0700
Subject: [PATCH 10/26] Fix signature in docstring

---
 Python/bltinmodule.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 9c107713f948da..b7eb0ef053c1cc 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2707,7 +2707,7 @@ static PyMethodDef zip_methods[] = {
 };
 
 PyDoc_STRVAR(zip_doc,
-"zip(*iterables) --> A zip object yielding tuples until an input is exhausted.\n\
+"zip(*iterables, strict=False) --> Yield tuples until an input is exhausted.\n\
 \n\
    >>> list(zip('abcdefg', range(3), range(4)))\n\
    [('a', 0, 0), ('b', 1, 1), ('c', 2, 2)]\n\

From 8e9b4e669af6dc0322fb8edffa9ff4187dd36265 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 10:12:12 -0700
Subject: [PATCH 11/26] Reword error messages

---
 Python/bltinmodule.c | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index b7eb0ef053c1cc..6b48ed4d728271 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2658,8 +2658,9 @@ zip_next(zipobject *lz)
     }
     if (i) {
         PyErr_Clear();
-        return PyErr_Format(PyExc_ValueError, "%s() argument %d is too short",
-                            Py_TYPE(lz)->tp_name, i + 1);
+        return PyErr_Format(PyExc_ValueError,
+                            "%s() arg %d is shorter than the preceding arg%s",
+                            Py_TYPE(lz)->tp_name, i + 1, i == 1 ? "" : "s");
     }
     for (i = 1; i < tuplesize; i++) {
         it = PyTuple_GET_ITEM(lz->ittuple, i);
@@ -2667,8 +2668,9 @@ zip_next(zipobject *lz)
         if (item) {
             Py_DECREF(item);
             PyErr_Clear();
-            return PyErr_Format(PyExc_ValueError, "%s() argument %d is too long",
-                                Py_TYPE(lz)->tp_name, i + 1);
+            return PyErr_Format(PyExc_ValueError,
+                                "%s() arg %d is longer than the preceding arg%s",
+                                Py_TYPE(lz)->tp_name, i + 1, i == 1 ? "" : "s");
         }
         if (PyErr_Occurred()) {
             return NULL;

From f2f3e03cfcbd7205ee1aac71a430736d44d15a80 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 10:27:47 -0700
Subject: [PATCH 12/26] blurb add

---
 .../2020-06-17-10-27-17.bpo-40636.MYaCIe.rst                  | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst

diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst b/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst
new file mode 100644
index 00000000000000..45714815de2e5e
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst	
@@ -0,0 +1,4 @@
+:func:`zip` now supports :pep:`618`'s ``strict`` parameter, which raises a
+:exc:`ValueError` if the arguments are exhausted at differing lengths.
+
+Patch by Brandt Bucher.

From 3ad4c925710d30ce9b3234858193bf3195088632 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 10:44:05 -0700
Subject: [PATCH 13/26] Improve error messages

---
 Python/bltinmodule.c | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 6b48ed4d728271..00332c5163e86d 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2658,9 +2658,10 @@ zip_next(zipobject *lz)
     }
     if (i) {
         PyErr_Clear();
+        const char* plural = i == 1 ? " " : "s 1-";
         return PyErr_Format(PyExc_ValueError,
-                            "%s() arg %d is shorter than the preceding arg%s",
-                            Py_TYPE(lz)->tp_name, i + 1, i == 1 ? "" : "s");
+                            "%s() argument %d is shorter than argument%s%d",
+                            Py_TYPE(lz)->tp_name, i + 1, plural, i);
     }
     for (i = 1; i < tuplesize; i++) {
         it = PyTuple_GET_ITEM(lz->ittuple, i);
@@ -2668,9 +2669,10 @@ zip_next(zipobject *lz)
         if (item) {
             Py_DECREF(item);
             PyErr_Clear();
+            const char* plural = i == 1 ? " " : "s 1-";
             return PyErr_Format(PyExc_ValueError,
-                                "%s() arg %d is longer than the preceding arg%s",
-                                Py_TYPE(lz)->tp_name, i + 1, i == 1 ? "" : "s");
+                                "%s() argument %d is longer than argument%s%d",
+                                Py_TYPE(lz)->tp_name, i + 1, plural, i);
         }
         if (PyErr_Occurred()) {
             return NULL;

From 0fa507ce9fbe08b67b0dfbeab73600d7be4ceccb Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 17:02:34 -0700
Subject: [PATCH 14/26] Fix NEWS

---
 .../Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst   | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst b/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst
index 45714815de2e5e..3772604327daf9 100644
--- a/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst	
+++ b/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst	
@@ -1,4 +1,3 @@
 :func:`zip` now supports :pep:`618`'s ``strict`` parameter, which raises a
 :exc:`ValueError` if the arguments are exhausted at differing lengths.
-
 Patch by Brandt Bucher.

From dca9af089ecf3be205baeb3239fca3d0c65c2dff Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 17:03:47 -0700
Subject: [PATCH 15/26] Add strict member to zipobject

---
 Python/bltinmodule.c | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 00332c5163e86d..bca32597408e6a 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2517,9 +2517,10 @@ builtin_issubclass_impl(PyObject *module, PyObject *cls,
 
 typedef struct {
     PyObject_HEAD
-    Py_ssize_t tuplesize;  /* negative when strict=True */
+    Py_ssize_t tuplesize;
     PyObject *ittuple;     /* tuple of iterators */
     PyObject *result;
+    int strict;
 } zipobject;
 
 static PyObject *
@@ -2583,8 +2584,9 @@ zip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         return NULL;
     }
     lz->ittuple = ittuple;
-    lz->tuplesize = strict ? -tuplesize : tuplesize;
+    lz->tuplesize = tuplesize;
     lz->result = result;
+    lz->strict = strict;
 
     return (PyObject *)lz;
 }
@@ -2610,7 +2612,7 @@ static PyObject *
 zip_next(zipobject *lz)
 {
     Py_ssize_t i;
-    Py_ssize_t tuplesize = Py_ABS(lz->tuplesize);
+    Py_ssize_t tuplesize = lz->tuplesize;
     PyObject *result = lz->result;
     PyObject *it;
     PyObject *item;
@@ -2625,7 +2627,7 @@ zip_next(zipobject *lz)
             item = (*Py_TYPE(it)->tp_iternext)(it);
             if (item == NULL) {
                 Py_DECREF(result);
-                if (lz->tuplesize < 0) {
+                if (lz->strict) {
                     goto check;
                 }
                 return NULL;
@@ -2643,7 +2645,7 @@ zip_next(zipobject *lz)
             item = (*Py_TYPE(it)->tp_iternext)(it);
             if (item == NULL) {
                 Py_DECREF(result);
-                if (lz->tuplesize < 0) {
+                if (lz->strict) {
                     goto check;
                 }
                 return NULL;
@@ -2685,7 +2687,7 @@ static PyObject *
 zip_reduce(zipobject *lz, PyObject *Py_UNUSED(ignored))
 {
     /* Just recreate the zip with the internal iterator tuple */
-    if (lz->tuplesize < 0) {
+    if (lz->strict) {
         return Py_BuildValue("OOO", Py_TYPE(lz), lz->ittuple, Py_True);
     }
     return Py_BuildValue("OO", Py_TYPE(lz), lz->ittuple);
@@ -2700,7 +2702,7 @@ zip_setstate(zipobject *lz, PyObject *state)
     if (strict < 0) {
         return NULL;
     }
-    lz->tuplesize = strict ? -Py_ABS(lz->tuplesize) : Py_ABS(lz->tuplesize);
+    lz->strict = strict;
     Py_RETURN_NONE;
 }
 

From aa5c8f647c20c250a79b94f9293caf9375678b60 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 17:12:12 -0700
Subject: [PATCH 16/26] Simplify error message construction

---
 Python/bltinmodule.c | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index bca32597408e6a..816851b9877a8f 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2662,8 +2662,8 @@ zip_next(zipobject *lz)
         PyErr_Clear();
         const char* plural = i == 1 ? " " : "s 1-";
         return PyErr_Format(PyExc_ValueError,
-                            "%s() argument %d is shorter than argument%s%d",
-                            Py_TYPE(lz)->tp_name, i + 1, plural, i);
+                            "zip() argument %d is shorter than argument%s%d",
+                            i + 1, plural, i);
     }
     for (i = 1; i < tuplesize; i++) {
         it = PyTuple_GET_ITEM(lz->ittuple, i);
@@ -2673,8 +2673,8 @@ zip_next(zipobject *lz)
             PyErr_Clear();
             const char* plural = i == 1 ? " " : "s 1-";
             return PyErr_Format(PyExc_ValueError,
-                                "%s() argument %d is longer than argument%s%d",
-                                Py_TYPE(lz)->tp_name, i + 1, plural, i);
+                                "zip() argument %d is longer than argument%s%d",
+                                i + 1, plural, i);
         }
         if (PyErr_Occurred()) {
             return NULL;

From 7d24e128f4da59f2598e42fccc79c8b2f503508c Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 17:15:28 -0700
Subject: [PATCH 17/26] Use PyTuple_Pack for pickling

---
 Python/bltinmodule.c | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 816851b9877a8f..d95bbae1672fec 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2688,9 +2688,9 @@ zip_reduce(zipobject *lz, PyObject *Py_UNUSED(ignored))
 {
     /* Just recreate the zip with the internal iterator tuple */
     if (lz->strict) {
-        return Py_BuildValue("OOO", Py_TYPE(lz), lz->ittuple, Py_True);
+        return PyTuple_Pack(3, Py_TYPE(lz), lz->ittuple, Py_True);
     }
-    return Py_BuildValue("OO", Py_TYPE(lz), lz->ittuple);
+    return PyTuple_Pack(2, Py_TYPE(lz), lz->ittuple);
 }
 
 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");

From 05183d916e28cb5384874e41c0a1e766c1aef153 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 17:25:09 -0700
Subject: [PATCH 18/26] Add round-trip pickle tests

---
 Lib/test/test_builtin.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index a6d2388ab03345..a044cd6388af28 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1573,6 +1573,28 @@ def test_zip_pickle(self):
             z1 = zip(a, b)
             self.check_iter_pickle(z1, t, proto)
 
+    def test_zip_pickle_strict(self):
+        a = (1, 2, 3)
+        b = (4, 5, 6)
+        t = [(1, 4), (2, 5), (3, 6)]
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            z1 = zip(a, b, strict=True)
+            self.check_iter_pickle(z1, t, proto)
+
+    def test_zip_pickle_strict_fail(self):
+        a = (1, 2, 3)
+        b = (4, 5, 6, 7)
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            z1 = zip(a, b, strict=True)
+            z2 = pickle.loads(pickle.dumps(z1, proto))
+            self.assertEqual(next(z1), next(z2))
+            self.assertEqual(next(z1), next(z2))
+            self.assertEqual(next(z1), next(z2))
+            with self.assertRaises(ValueError):
+                next(z1)
+            with self.assertRaises(ValueError):
+                next(z2)
+
     def test_zip_bad_iterable(self):
         exception = TypeError()
 

From 28edf45b8ac7ec36085ee81bdfe9c51cedea99a0 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Wed, 17 Jun 2020 21:19:26 -0700
Subject: [PATCH 19/26] Add tests for pickle stability

---
 Lib/test/test_builtin.py | 43 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index a044cd6388af28..4ada49790714fc 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1595,6 +1595,49 @@ def test_zip_pickle_strict_fail(self):
             with self.assertRaises(ValueError):
                 next(z2)
 
+    def test_zip_pickle_stability(self):
+        # Pickles of zip((1, 2, 3), (4, 5, 6)) dumped from 3.9:
+        pickles = [
+            b'citertools\nizip\np0\n(c__builtin__\niter\np1\n((I1\nI2\nI3\ntp2\ntp3\nRp4\nI0\nbg1\n((I4\nI5\nI6\ntp5\ntp6\nRp7\nI0\nbtp8\nRp9\n.',
+            b'citertools\nizip\nq\x00(c__builtin__\niter\nq\x01((K\x01K\x02K\x03tq\x02tq\x03Rq\x04K\x00bh\x01((K\x04K\x05K\x06tq\x05tq\x06Rq\x07K\x00btq\x08Rq\t.',
+            b'\x80\x02citertools\nizip\nq\x00c__builtin__\niter\nq\x01K\x01K\x02K\x03\x87q\x02\x85q\x03Rq\x04K\x00bh\x01K\x04K\x05K\x06\x87q\x05\x85q\x06Rq\x07K\x00b\x86q\x08Rq\t.',
+            b'\x80\x03cbuiltins\nzip\nq\x00cbuiltins\niter\nq\x01K\x01K\x02K\x03\x87q\x02\x85q\x03Rq\x04K\x00bh\x01K\x04K\x05K\x06\x87q\x05\x85q\x06Rq\x07K\x00b\x86q\x08Rq\t.',
+            b'\x80\x04\x95L\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x03zip\x94\x93\x94\x8c\x08builtins\x94\x8c\x04iter\x94\x93\x94K\x01K\x02K\x03\x87\x94\x85\x94R\x94K\x00bh\x05K\x04K\x05K\x06\x87\x94\x85\x94R\x94K\x00b\x86\x94R\x94.',
+            b'\x80\x05\x95L\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x03zip\x94\x93\x94\x8c\x08builtins\x94\x8c\x04iter\x94\x93\x94K\x01K\x02K\x03\x87\x94\x85\x94R\x94K\x00bh\x05K\x04K\x05K\x06\x87\x94\x85\x94R\x94K\x00b\x86\x94R\x94.',
+        ]
+        for protocol, dump in enumerate(pickles):
+            z1 = zip((1, 2, 3), (4, 5, 6))
+            z2 = zip((1, 2, 3), (4, 5, 6), strict=False)
+            z3 = pickle.loads(dump)
+            l3 = list(z3)
+            self.assertEqual(type(z3), zip)
+            self.assertEqual(pickle.dumps(z1, protocol), dump)
+            self.assertEqual(pickle.dumps(z2, protocol), dump)
+            self.assertEqual(list(z1), l3)
+            self.assertEqual(list(z2), l3)
+
+    def test_zip_pickle_strict_stability(self):
+        # Pickles of zip((1, 2, 3), (4, 5), strict=True) dumped from 3.10:
+        pickles = [
+            b'citertools\nizip\np0\n(c__builtin__\niter\np1\n((I1\nI2\nI3\ntp2\ntp3\nRp4\nI0\nbg1\n((I4\nI5\ntp5\ntp6\nRp7\nI0\nbtp8\nRp9\nI01\nb.',
+            b'citertools\nizip\nq\x00(c__builtin__\niter\nq\x01((K\x01K\x02K\x03tq\x02tq\x03Rq\x04K\x00bh\x01((K\x04K\x05tq\x05tq\x06Rq\x07K\x00btq\x08Rq\tI01\nb.',
+            b'\x80\x02citertools\nizip\nq\x00c__builtin__\niter\nq\x01K\x01K\x02K\x03\x87q\x02\x85q\x03Rq\x04K\x00bh\x01K\x04K\x05\x86q\x05\x85q\x06Rq\x07K\x00b\x86q\x08Rq\t\x88b.',
+            b'\x80\x03cbuiltins\nzip\nq\x00cbuiltins\niter\nq\x01K\x01K\x02K\x03\x87q\x02\x85q\x03Rq\x04K\x00bh\x01K\x04K\x05\x86q\x05\x85q\x06Rq\x07K\x00b\x86q\x08Rq\t\x88b.',
+            b'\x80\x04\x95L\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x03zip\x94\x93\x94\x8c\x08builtins\x94\x8c\x04iter\x94\x93\x94K\x01K\x02K\x03\x87\x94\x85\x94R\x94K\x00bh\x05K\x04K\x05\x86\x94\x85\x94R\x94K\x00b\x86\x94R\x94\x88b.',
+            b'\x80\x05\x95L\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x03zip\x94\x93\x94\x8c\x08builtins\x94\x8c\x04iter\x94\x93\x94K\x01K\x02K\x03\x87\x94\x85\x94R\x94K\x00bh\x05K\x04K\x05\x86\x94\x85\x94R\x94K\x00b\x86\x94R\x94\x88b.',
+        ]
+        for protocol, dump in enumerate(pickles):
+            z1 = zip((1, 2, 3), (4, 5), strict=True)
+            z2 = pickle.loads(dump)
+            self.assertEqual(pickle.dumps(z1, protocol), dump)
+            self.assertEqual(type(z1), zip)
+            self.assertEqual(next(z1), next(z2))
+            self.assertEqual(next(z1), next(z2))
+            with self.assertRaises(ValueError):
+                next(z1)
+            with self.assertRaises(ValueError):
+                next(z2)
+
     def test_zip_bad_iterable(self):
         exception = TypeError()
 

From baff9c81ea9a05ef84eaa4d280e762c2ecfca27e Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 18 Jun 2020 10:16:36 -0700
Subject: [PATCH 20/26] Add comments and improve error handling.

---
 Python/bltinmodule.c | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index d95bbae1672fec..a7d013095feeec 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2655,12 +2655,17 @@ zip_next(zipobject *lz)
     }
     return result;
 check:
-    if (PyErr_Occurred() && !PyErr_ExceptionMatches(PyExc_StopIteration)) {
-        return NULL;
+    if (PyErr_Occurred()) {
+        if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
+            // next() on argument i raised an exception (not StopIteration)
+            return NULL;
+        }
+        PyErr_Clear();
     }
     if (i) {
-        PyErr_Clear();
-        const char* plural = i == 1 ? " " : "s 1-";
+        // ValueError: zip() argument 2 is shorter than argument 1
+        // ValueError: zip() argument 3 is shorter than arguments 1-2
+        const char* plural = i == 1 ? " " : "s 1-";  //         ^^^^
         return PyErr_Format(PyExc_ValueError,
                             "zip() argument %d is shorter than argument%s%d",
                             i + 1, plural, i);
@@ -2670,16 +2675,21 @@ zip_next(zipobject *lz)
         item = (*Py_TYPE(it)->tp_iternext)(it);
         if (item) {
             Py_DECREF(item);
-            PyErr_Clear();
             const char* plural = i == 1 ? " " : "s 1-";
             return PyErr_Format(PyExc_ValueError,
                                 "zip() argument %d is longer than argument%s%d",
                                 i + 1, plural, i);
         }
         if (PyErr_Occurred()) {
-            return NULL;
+            if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
+                // next() on argument i raised an exception (not StopIteration)
+                return NULL;
+            }
+            PyErr_Clear();
         }
+        // Argument i is exhausted. So far so good...
     }
+    // All arguments are exhausted. Success!
     return NULL;
 }
 

From 2ac303679648caab1f75146d0a0adc8988207ec5 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 18 Jun 2020 10:23:13 -0700
Subject: [PATCH 21/26] Add tests for error handling

---
 Lib/test/test_builtin.py | 58 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index 4ada49790714fc..1d2d5ca9e41f3c 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1669,6 +1669,64 @@ def test_zip_strict_iterators(self):
         self.assertEqual(next(x), 2)
         self.assertEqual(next(z), 1)
 
+    def test_zip_strict_error_handling(self):
+        class E(Exception):
+            pass
+        class I:
+            def __init__(self, size):
+                self.size = size
+            def __iter__(self):
+                return self
+            def __next__(self):
+                self.size -= 1
+                if self.size < 0:
+                    raise E
+                return self.size
+        with self.assertRaises(E):
+            list(zip("AB", I(1), strict=True))
+        with self.assertRaises(ValueError):
+            list(zip("AB", I(2), "A", strict=True))
+        with self.assertRaises(E):
+            list(zip("AB", I(2), "ABC", strict=True))
+        with self.assertRaises(ValueError):
+            list(zip("AB", I(3), strict=True))
+        with self.assertRaises(E):
+            list(zip(I(1), "AB", strict=True))
+        with self.assertRaises(ValueError):
+            list(zip(I(2), "A", strict=True))
+        with self.assertRaises(E):
+            list(zip(I(2), "ABC", strict=True))
+        with self.assertRaises(ValueError):
+            list(zip(I(3), "AB", strict=True))
+
+    def test_zip_strict_error_handling_stopiteration(self):
+        class I:
+            def __init__(self, size):
+                self.size = size
+            def __iter__(self):
+                return self
+            def __next__(self):
+                self.size -= 1
+                if self.size < 0:
+                    raise StopIteration
+                return self.size
+        with self.assertRaises(ValueError):
+            list(zip("AB", I(1), strict=True))
+        with self.assertRaises(ValueError):
+            list(zip("AB", I(2), "A", strict=True))
+        with self.assertRaises(ValueError):
+            list(zip("AB", I(2), "ABC", strict=True))
+        with self.assertRaises(ValueError):
+            list(zip("AB", I(3), strict=True))
+        with self.assertRaises(ValueError):
+            list(zip(I(1), "AB", strict=True))
+        with self.assertRaises(ValueError):
+            list(zip(I(2), "A", strict=True))
+        with self.assertRaises(ValueError):
+            list(zip(I(2), "ABC", strict=True))
+        with self.assertRaises(ValueError):
+            list(zip(I(3), "AB", strict=True))
+
     def test_format(self):
         # Test the basic machinery of the format() builtin.  Don't test
         #  the specifics of the various formatters

From ea2a7cfb090b306b0f80d40aa9ff398a7a7b5963 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 18 Jun 2020 10:23:22 -0700
Subject: [PATCH 22/26] Reword NEWS

---
 .../Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst b/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst
index 3772604327daf9..ba26ad9373ce38 100644
--- a/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst	
+++ b/Misc/NEWS.d/next/Core and Builtins/2020-06-17-10-27-17.bpo-40636.MYaCIe.rst	
@@ -1,3 +1,3 @@
 :func:`zip` now supports :pep:`618`'s ``strict`` parameter, which raises a
-:exc:`ValueError` if the arguments are exhausted at differing lengths.
+:exc:`ValueError` if the arguments are exhausted at different lengths.
 Patch by Brandt Bucher.

From 3613e34202916e7e5e970c23648f03b7dc8131d3 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 18 Jun 2020 11:05:27 -0700
Subject: [PATCH 23/26] Refactor tests

---
 Lib/test/test_builtin.py | 97 +++++++++++++++++++++-------------------
 1 file changed, 50 insertions(+), 47 deletions(-)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index 1d2d5ca9e41f3c..5f6e33bd57a8aa 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1521,6 +1521,14 @@ def test_vars(self):
         self.assertRaises(TypeError, vars, 42)
         self.assertEqual(vars(self.C_get_vars()), {'a':2})
 
+    def iter_error(self, iterable, error):
+        """Collect `iterable` into a list, catching an expected `error`."""
+        items = []
+        with self.assertRaises(error):
+            for item in iterable:
+                items.append(item)
+        return items
+
     def test_zip(self):
         a = (1, 2, 3)
         b = (4, 5, 6)
@@ -1584,16 +1592,12 @@ def test_zip_pickle_strict(self):
     def test_zip_pickle_strict_fail(self):
         a = (1, 2, 3)
         b = (4, 5, 6, 7)
+        t = [(1, 4), (2, 5), (3, 6)]
         for proto in range(pickle.HIGHEST_PROTOCOL + 1):
             z1 = zip(a, b, strict=True)
             z2 = pickle.loads(pickle.dumps(z1, proto))
-            self.assertEqual(next(z1), next(z2))
-            self.assertEqual(next(z1), next(z2))
-            self.assertEqual(next(z1), next(z2))
-            with self.assertRaises(ValueError):
-                next(z1)
-            with self.assertRaises(ValueError):
-                next(z2)
+            self.assertEqual(self.iter_error(z1, ValueError), t)
+            self.assertEqual(self.iter_error(z2, ValueError), t)
 
     def test_zip_pickle_stability(self):
         # Pickles of zip((1, 2, 3), (4, 5, 6)) dumped from 3.9:
@@ -1626,17 +1630,16 @@ def test_zip_pickle_strict_stability(self):
             b'\x80\x04\x95L\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x03zip\x94\x93\x94\x8c\x08builtins\x94\x8c\x04iter\x94\x93\x94K\x01K\x02K\x03\x87\x94\x85\x94R\x94K\x00bh\x05K\x04K\x05\x86\x94\x85\x94R\x94K\x00b\x86\x94R\x94\x88b.',
             b'\x80\x05\x95L\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x03zip\x94\x93\x94\x8c\x08builtins\x94\x8c\x04iter\x94\x93\x94K\x01K\x02K\x03\x87\x94\x85\x94R\x94K\x00bh\x05K\x04K\x05\x86\x94\x85\x94R\x94K\x00b\x86\x94R\x94\x88b.',
         ]
+        a = (1, 2, 3)
+        b = (4, 5)
+        t = [(1, 4), (2, 5)]
         for protocol, dump in enumerate(pickles):
-            z1 = zip((1, 2, 3), (4, 5), strict=True)
+            z1 = zip(a, b, strict=True)
             z2 = pickle.loads(dump)
             self.assertEqual(pickle.dumps(z1, protocol), dump)
-            self.assertEqual(type(z1), zip)
-            self.assertEqual(next(z1), next(z2))
-            self.assertEqual(next(z1), next(z2))
-            with self.assertRaises(ValueError):
-                next(z1)
-            with self.assertRaises(ValueError):
-                next(z2)
+            self.assertEqual(type(z2), zip)
+            self.assertEqual(self.iter_error(z1, ValueError), t)
+            self.assertEqual(self.iter_error(z2, ValueError), t)
 
     def test_zip_bad_iterable(self):
         exception = TypeError()
@@ -1682,22 +1685,22 @@ def __next__(self):
                 if self.size < 0:
                     raise E
                 return self.size
-        with self.assertRaises(E):
-            list(zip("AB", I(1), strict=True))
-        with self.assertRaises(ValueError):
-            list(zip("AB", I(2), "A", strict=True))
-        with self.assertRaises(E):
-            list(zip("AB", I(2), "ABC", strict=True))
-        with self.assertRaises(ValueError):
-            list(zip("AB", I(3), strict=True))
-        with self.assertRaises(E):
-            list(zip(I(1), "AB", strict=True))
-        with self.assertRaises(ValueError):
-            list(zip(I(2), "A", strict=True))
-        with self.assertRaises(E):
-            list(zip(I(2), "ABC", strict=True))
-        with self.assertRaises(ValueError):
-            list(zip(I(3), "AB", strict=True))
+        l1 = self.iter_error(zip("AB", I(1), strict=True), E)
+        self.assertEqual(l1, [("A", 0)])
+        l2 = self.iter_error(zip("AB", I(2), "A", strict=True), ValueError)
+        self.assertEqual(l2, [("A", 1, "A")])
+        l3 = self.iter_error(zip("AB", I(2), "ABC", strict=True), E)
+        self.assertEqual(l3, [("A", 1, "A"), ("B", 0, "B")])
+        l4 = self.iter_error(zip("AB", I(3), strict=True), ValueError)
+        self.assertEqual(l4, [("A", 2), ("B", 1)])
+        l5 = self.iter_error(zip(I(1), "AB", strict=True), E)
+        self.assertEqual(l5, [(0, "A")])
+        l6 = self.iter_error(zip(I(2), "A", strict=True), ValueError)
+        self.assertEqual(l6, [(1, "A")])
+        l7 = self.iter_error(zip(I(2), "ABC", strict=True), E)
+        self.assertEqual(l7, [(1, "A"), (0, "B")])
+        l8 = self.iter_error(zip(I(3), "AB", strict=True), ValueError)
+        self.assertEqual(l8, [(2, "A"), (1, "B")])
 
     def test_zip_strict_error_handling_stopiteration(self):
         class I:
@@ -1710,22 +1713,22 @@ def __next__(self):
                 if self.size < 0:
                     raise StopIteration
                 return self.size
-        with self.assertRaises(ValueError):
-            list(zip("AB", I(1), strict=True))
-        with self.assertRaises(ValueError):
-            list(zip("AB", I(2), "A", strict=True))
-        with self.assertRaises(ValueError):
-            list(zip("AB", I(2), "ABC", strict=True))
-        with self.assertRaises(ValueError):
-            list(zip("AB", I(3), strict=True))
-        with self.assertRaises(ValueError):
-            list(zip(I(1), "AB", strict=True))
-        with self.assertRaises(ValueError):
-            list(zip(I(2), "A", strict=True))
-        with self.assertRaises(ValueError):
-            list(zip(I(2), "ABC", strict=True))
-        with self.assertRaises(ValueError):
-            list(zip(I(3), "AB", strict=True))
+        l1 = self.iter_error(zip("AB", I(1), strict=True), ValueError)
+        self.assertEqual(l1, [("A", 0)])
+        l2 = self.iter_error(zip("AB", I(2), "A", strict=True), ValueError)
+        self.assertEqual(l2, [("A", 1, "A")])
+        l3 = self.iter_error(zip("AB", I(2), "ABC", strict=True), ValueError)
+        self.assertEqual(l3, [("A", 1, "A"), ("B", 0, "B")])
+        l4 = self.iter_error(zip("AB", I(3), strict=True), ValueError)
+        self.assertEqual(l4, [("A", 2), ("B", 1)])
+        l5 = self.iter_error(zip(I(1), "AB", strict=True), ValueError)
+        self.assertEqual(l5, [(0, "A")])
+        l6 = self.iter_error(zip(I(2), "A", strict=True), ValueError)
+        self.assertEqual(l6, [(1, "A")])
+        l7 = self.iter_error(zip(I(2), "ABC", strict=True), ValueError)
+        self.assertEqual(l7, [(1, "A"), (0, "B")])
+        l8 = self.iter_error(zip(I(3), "AB", strict=True), ValueError)
+        self.assertEqual(l8, [(2, "A"), (1, "B")])
 
     def test_format(self):
         # Test the basic machinery of the format() builtin.  Don't test

From 329afd924396db09c8778c7acbf39213e9b7fdf3 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 18 Jun 2020 11:10:36 -0700
Subject: [PATCH 24/26] Clean up tests

---
 Lib/test/test_builtin.py | 33 +++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index 5f6e33bd57a8aa..ac59b47acbdc70 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1673,8 +1673,10 @@ def test_zip_strict_iterators(self):
         self.assertEqual(next(z), 1)
 
     def test_zip_strict_error_handling(self):
+
         class E(Exception):
             pass
+
         class I:
             def __init__(self, size):
                 self.size = size
@@ -1685,24 +1687,26 @@ def __next__(self):
                 if self.size < 0:
                     raise E
                 return self.size
+
         l1 = self.iter_error(zip("AB", I(1), strict=True), E)
-        self.assertEqual(l1, [("A", 0)])
         l2 = self.iter_error(zip("AB", I(2), "A", strict=True), ValueError)
-        self.assertEqual(l2, [("A", 1, "A")])
         l3 = self.iter_error(zip("AB", I(2), "ABC", strict=True), E)
-        self.assertEqual(l3, [("A", 1, "A"), ("B", 0, "B")])
         l4 = self.iter_error(zip("AB", I(3), strict=True), ValueError)
-        self.assertEqual(l4, [("A", 2), ("B", 1)])
         l5 = self.iter_error(zip(I(1), "AB", strict=True), E)
-        self.assertEqual(l5, [(0, "A")])
         l6 = self.iter_error(zip(I(2), "A", strict=True), ValueError)
-        self.assertEqual(l6, [(1, "A")])
         l7 = self.iter_error(zip(I(2), "ABC", strict=True), E)
-        self.assertEqual(l7, [(1, "A"), (0, "B")])
         l8 = self.iter_error(zip(I(3), "AB", strict=True), ValueError)
+        self.assertEqual(l1, [("A", 0)])
+        self.assertEqual(l2, [("A", 1, "A")])
+        self.assertEqual(l3, [("A", 1, "A"), ("B", 0, "B")])
+        self.assertEqual(l4, [("A", 2), ("B", 1)])
+        self.assertEqual(l5, [(0, "A")])
+        self.assertEqual(l6, [(1, "A")])
+        self.assertEqual(l7, [(1, "A"), (0, "B")])
         self.assertEqual(l8, [(2, "A"), (1, "B")])
 
     def test_zip_strict_error_handling_stopiteration(self):
+
         class I:
             def __init__(self, size):
                 self.size = size
@@ -1713,21 +1717,22 @@ def __next__(self):
                 if self.size < 0:
                     raise StopIteration
                 return self.size
+
         l1 = self.iter_error(zip("AB", I(1), strict=True), ValueError)
-        self.assertEqual(l1, [("A", 0)])
         l2 = self.iter_error(zip("AB", I(2), "A", strict=True), ValueError)
-        self.assertEqual(l2, [("A", 1, "A")])
         l3 = self.iter_error(zip("AB", I(2), "ABC", strict=True), ValueError)
-        self.assertEqual(l3, [("A", 1, "A"), ("B", 0, "B")])
         l4 = self.iter_error(zip("AB", I(3), strict=True), ValueError)
-        self.assertEqual(l4, [("A", 2), ("B", 1)])
         l5 = self.iter_error(zip(I(1), "AB", strict=True), ValueError)
-        self.assertEqual(l5, [(0, "A")])
         l6 = self.iter_error(zip(I(2), "A", strict=True), ValueError)
-        self.assertEqual(l6, [(1, "A")])
         l7 = self.iter_error(zip(I(2), "ABC", strict=True), ValueError)
-        self.assertEqual(l7, [(1, "A"), (0, "B")])
         l8 = self.iter_error(zip(I(3), "AB", strict=True), ValueError)
+        self.assertEqual(l1, [("A", 0)])
+        self.assertEqual(l2, [("A", 1, "A")])
+        self.assertEqual(l3, [("A", 1, "A"), ("B", 0, "B")])
+        self.assertEqual(l4, [("A", 2), ("B", 1)])
+        self.assertEqual(l5, [(0, "A")])
+        self.assertEqual(l6, [(1, "A")])
+        self.assertEqual(l7, [(1, "A"), (0, "B")])
         self.assertEqual(l8, [(2, "A"), (1, "B")])
 
     def test_format(self):

From e86daab3d36c056082c81023c99fbbc75bd80b35 Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 18 Jun 2020 15:56:19 -0700
Subject: [PATCH 25/26] Remove comment

---
 Python/bltinmodule.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index a7d013095feeec..005368c4c875d7 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -2665,7 +2665,7 @@ zip_next(zipobject *lz)
     if (i) {
         // ValueError: zip() argument 2 is shorter than argument 1
         // ValueError: zip() argument 3 is shorter than arguments 1-2
-        const char* plural = i == 1 ? " " : "s 1-";  //         ^^^^
+        const char* plural = i == 1 ? " " : "s 1-";
         return PyErr_Format(PyExc_ValueError,
                             "zip() argument %d is shorter than argument%s%d",
                             i + 1, plural, i);

From d77134e92b042b18c45dbaadd72d9593199cc5fa Mon Sep 17 00:00:00 2001
From: Brandt Bucher <brandtbucher@gmail.com>
Date: Thu, 18 Jun 2020 15:56:33 -0700
Subject: [PATCH 26/26] Rename/reorder tests

---
 Lib/test/test_builtin.py | 40 ++++++++++++++++++++--------------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index ac59b47acbdc70..40df7b606ae5e0 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1674,10 +1674,10 @@ def test_zip_strict_iterators(self):
 
     def test_zip_strict_error_handling(self):
 
-        class E(Exception):
+        class Error(Exception):
             pass
 
-        class I:
+        class Iter:
             def __init__(self, size):
                 self.size = size
             def __iter__(self):
@@ -1685,29 +1685,29 @@ def __iter__(self):
             def __next__(self):
                 self.size -= 1
                 if self.size < 0:
-                    raise E
+                    raise Error
                 return self.size
 
-        l1 = self.iter_error(zip("AB", I(1), strict=True), E)
-        l2 = self.iter_error(zip("AB", I(2), "A", strict=True), ValueError)
-        l3 = self.iter_error(zip("AB", I(2), "ABC", strict=True), E)
-        l4 = self.iter_error(zip("AB", I(3), strict=True), ValueError)
-        l5 = self.iter_error(zip(I(1), "AB", strict=True), E)
-        l6 = self.iter_error(zip(I(2), "A", strict=True), ValueError)
-        l7 = self.iter_error(zip(I(2), "ABC", strict=True), E)
-        l8 = self.iter_error(zip(I(3), "AB", strict=True), ValueError)
+        l1 = self.iter_error(zip("AB", Iter(1), strict=True), Error)
         self.assertEqual(l1, [("A", 0)])
+        l2 = self.iter_error(zip("AB", Iter(2), "A", strict=True), ValueError)
         self.assertEqual(l2, [("A", 1, "A")])
+        l3 = self.iter_error(zip("AB", Iter(2), "ABC", strict=True), Error)
         self.assertEqual(l3, [("A", 1, "A"), ("B", 0, "B")])
+        l4 = self.iter_error(zip("AB", Iter(3), strict=True), ValueError)
         self.assertEqual(l4, [("A", 2), ("B", 1)])
+        l5 = self.iter_error(zip(Iter(1), "AB", strict=True), Error)
         self.assertEqual(l5, [(0, "A")])
+        l6 = self.iter_error(zip(Iter(2), "A", strict=True), ValueError)
         self.assertEqual(l6, [(1, "A")])
+        l7 = self.iter_error(zip(Iter(2), "ABC", strict=True), Error)
         self.assertEqual(l7, [(1, "A"), (0, "B")])
+        l8 = self.iter_error(zip(Iter(3), "AB", strict=True), ValueError)
         self.assertEqual(l8, [(2, "A"), (1, "B")])
 
     def test_zip_strict_error_handling_stopiteration(self):
 
-        class I:
+        class Iter:
             def __init__(self, size):
                 self.size = size
             def __iter__(self):
@@ -1718,21 +1718,21 @@ def __next__(self):
                     raise StopIteration
                 return self.size
 
-        l1 = self.iter_error(zip("AB", I(1), strict=True), ValueError)
-        l2 = self.iter_error(zip("AB", I(2), "A", strict=True), ValueError)
-        l3 = self.iter_error(zip("AB", I(2), "ABC", strict=True), ValueError)
-        l4 = self.iter_error(zip("AB", I(3), strict=True), ValueError)
-        l5 = self.iter_error(zip(I(1), "AB", strict=True), ValueError)
-        l6 = self.iter_error(zip(I(2), "A", strict=True), ValueError)
-        l7 = self.iter_error(zip(I(2), "ABC", strict=True), ValueError)
-        l8 = self.iter_error(zip(I(3), "AB", strict=True), ValueError)
+        l1 = self.iter_error(zip("AB", Iter(1), strict=True), ValueError)
         self.assertEqual(l1, [("A", 0)])
+        l2 = self.iter_error(zip("AB", Iter(2), "A", strict=True), ValueError)
         self.assertEqual(l2, [("A", 1, "A")])
+        l3 = self.iter_error(zip("AB", Iter(2), "ABC", strict=True), ValueError)
         self.assertEqual(l3, [("A", 1, "A"), ("B", 0, "B")])
+        l4 = self.iter_error(zip("AB", Iter(3), strict=True), ValueError)
         self.assertEqual(l4, [("A", 2), ("B", 1)])
+        l5 = self.iter_error(zip(Iter(1), "AB", strict=True), ValueError)
         self.assertEqual(l5, [(0, "A")])
+        l6 = self.iter_error(zip(Iter(2), "A", strict=True), ValueError)
         self.assertEqual(l6, [(1, "A")])
+        l7 = self.iter_error(zip(Iter(2), "ABC", strict=True), ValueError)
         self.assertEqual(l7, [(1, "A"), (0, "B")])
+        l8 = self.iter_error(zip(Iter(3), "AB", strict=True), ValueError)
         self.assertEqual(l8, [(2, "A"), (1, "B")])
 
     def test_format(self):