bpo-40636: PEP 618 implementation

Lib/test/test_builtin.py

@@ -1585,6 +1585,25 @@ def __iter__(self):
 
         self.assertIs(cm.exception, exception)
 
+    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', strict=True))
+        self.assertRaises(ValueError, tuple,
+                          zip((1, 2), 'abc', strict=True))
+        self.assertRaises(ValueError, tuple,
+                          zip((1, 2), (1, 2), 'abc', strict=True))
+
+    def test_zip_strict_iterators(self):
+        x = iter(range(5))
+        y = [0]
+        z = iter(range(5))
+        self.assertRaises(ValueError, list,
+                          (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
         #  the specifics of the various formatters

Python/bltinmodule.c

@@ -2517,8 +2517,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 strict=True */
+    PyObject *ittuple;     /* tuple of iterators */
     PyObject *result;
 } zipobject;
 
@@ -2530,9 +2530,21 @@ zip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     PyObject *ittuple;  /* tuple of iterators */
     PyObject *result;
     Py_ssize_t tuplesize;
+    int strict = 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[] = {"strict", NULL};
+        int parsed = PyArg_ParseTupleAndKeywords(
+                empty, kwds, "|$p:zip", kwlist, &strict);
+        Py_DECREF(empty);
+        if (!parsed) {
+            return NULL;
+        }
+    }
 
     /* args must be a tuple */
     assert(PyTuple_Check(args));
@@ -2571,7 +2583,7 @@ zip_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         return NULL;
     }
     lz->ittuple = ittuple;
-    lz->tuplesize = tuplesize;
+    lz->tuplesize = strict ? -tuplesize : tuplesize;
     lz->result = result;
 
     return (PyObject *)lz;
@@ -2598,7 +2610,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;
@@ -2613,6 +2625,9 @@ zip_next(zipobject *lz)
             item = (*Py_TYPE(it)->tp_iternext)(it);
             if (item == NULL) {
                 Py_DECREF(result);
+                if (lz->tuplesize < 0) {
+                    goto check;
+                }
                 return NULL;
             }
             olditem = PyTuple_GET_ITEM(result, i);
@@ -2628,24 +2643,67 @@ zip_next(zipobject *lz)
             item = (*Py_TYPE(it)->tp_iternext)(it);
             if (item == NULL) {
                 Py_DECREF(result);
+                if (lz->tuplesize < 0) {
+                    goto check;
+                }
                 return NULL;
             }
             PyTuple_SET_ITEM(result, i, item);
         }
     }
     return result;
+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);
+    }
+    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);
+            PyErr_Clear();
+            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))
 {
     /* Just recreate the zip with the internal iterator tuple */
+    if (lz->tuplesize < 0) {
+        return Py_BuildValue("OOO", Py_TYPE(lz), lz->ittuple, Py_True);
+    }
     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[] = {
     {"__reduce__",   (PyCFunction)zip_reduce,   METH_NOARGS, reduce_doc},
-    {NULL,           NULL}           /* sentinel */
+    {"__setstate__", (PyCFunction)zip_setstate, METH_O,      setstate_doc},
+    {NULL}  /* sentinel */
 };
 
 PyDoc_STRVAR(zip_doc,
@@ -2657,7 +2715,10 @@ PyDoc_STRVAR(zip_doc,
 The zip object yields n-length tuples, where n is the number of iterables\n\
 passed as positional arguments to zip().  The i-th element in every tuple\n\
 comes from the i-th iterable argument to zip().  This continues until the\n\
-shortest argument is exhausted.");
+shortest argument is exhausted.\n\
+\n\
+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)