gh-119793: Add optional length-checking to map()

Lib/test/test_itertools.py

@@ -2140,10 +2140,10 @@ class subclass(cls):
                     subclass(*args, newarg=3)
 
         for cls, args, result in testcases:
-            # Constructors of repeat, zip, compress accept keyword arguments.
+            # Constructors of repeat, zip, map, compress accept keyword arguments.
             # Their subclasses need overriding __new__ to support new
             # keyword arguments.
-            if cls in [repeat, zip, compress]:
+            if cls in [repeat, zip, map, compress]:
                 continue
             with self.subTest(cls):
                 class subclass_with_init(cls):

Python/bltinmodule.c

@@ -1303,6 +1303,7 @@ typedef struct {
     PyObject_HEAD
     PyObject *iters;
     PyObject *func;
+    int strict;
 } mapobject;
 
 static PyObject *
@@ -1311,10 +1312,21 @@ map_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     PyObject *it, *iters, *func;
     mapobject *lz;
     Py_ssize_t numargs, i;
+    int strict = 0;
 
-    if ((type == &PyMap_Type || type->tp_init == PyMap_Type.tp_init) &&
-        !_PyArg_NoKeywords("map", 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:map", kwlist, &strict);
+        Py_DECREF(empty);
+        if (!parsed) {
+            return NULL;
+        }
+    }
 
     numargs = PyTuple_Size(args);
     if (numargs < 2) {
@@ -1346,47 +1358,7 @@ map_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     lz->iters = iters;
     func = PyTuple_GET_ITEM(args, 0);
     lz->func = Py_NewRef(func);
-
-    return (PyObject *)lz;
-}
-
-static PyObject *
-map_vectorcall(PyObject *type, PyObject * const*args,
-                size_t nargsf, PyObject *kwnames)
-{
-    PyTypeObject *tp = _PyType_CAST(type);
-    if (tp == &PyMap_Type && !_PyArg_NoKwnames("map", kwnames)) {
-        return NULL;
-    }
-
-    Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
-    if (nargs < 2) {
-        PyErr_SetString(PyExc_TypeError,
-           "map() must have at least two arguments.");
-        return NULL;
-    }
-
-    PyObject *iters = PyTuple_New(nargs-1);
-    if (iters == NULL) {
-        return NULL;
-    }
-
-    for (int i=1; i<nargs; i++) {
-        PyObject *it = PyObject_GetIter(args[i]);
-        if (it == NULL) {
-            Py_DECREF(iters);
-            return NULL;
-        }
-        PyTuple_SET_ITEM(iters, i-1, it);
-    }
-
-    mapobject *lz = (mapobject *)tp->tp_alloc(tp, 0);
-    if (lz == NULL) {
-        Py_DECREF(iters);
-        return NULL;
-    }
-    lz->iters = iters;
-    lz->func = Py_NewRef(args[0]);
+    lz->strict = strict;
 
     return (PyObject *)lz;
 }
@@ -1411,6 +1383,7 @@ map_traverse(mapobject *lz, visitproc visit, void *arg)
 static PyObject *
 map_next(mapobject *lz)
 {
+    Py_ssize_t i;
     PyObject *small_stack[_PY_FASTCALL_SMALL_STACK];
     PyObject **stack;
     PyObject *result = NULL;
@@ -1429,10 +1402,13 @@ map_next(mapobject *lz)
     }
 
     Py_ssize_t nargs = 0;
-    for (Py_ssize_t i=0; i < niters; i++) {
+    for (i=0; i < niters; i++) {
         PyObject *it = PyTuple_GET_ITEM(lz->iters, i);
         PyObject *val = Py_TYPE(it)->tp_iternext(it);
         if (val == NULL) {
+            if (lz->strict) {
+                goto check;
+            }
             goto exit;
         }
         stack[i] = val;
@@ -1449,6 +1425,43 @@ map_next(mapobject *lz)
         PyMem_Free(stack);
     }
     return result;
+check:
+    if (PyErr_Occurred()) {
+        if (!PyErr_ExceptionMatches(PyExc_StopIteration)) {
+            // next() on argument i raised an exception (not StopIteration)
+            return NULL;
+        }
+        PyErr_Clear();
+    }
+    if (i) {
+        // ValueError: map() argument 2 is shorter than argument 1
+        // ValueError: map() argument 3 is shorter than arguments 1-2
+        const char* plural = i == 1 ? " " : "s 1-";
+        return PyErr_Format(PyExc_ValueError,
+                            "map() argument %d is shorter than argument%s%d",
+                            i + 1, plural, i);
+    }
+    for (i = 1; i < niters; i++) {
+        PyObject *it = PyTuple_GET_ITEM(lz->iters, i);
+        PyObject *val = (*Py_TYPE(it)->tp_iternext)(it);
+        if (val) {
+            Py_DECREF(val);
+            const char* plural = i == 1 ? " " : "s 1-";
+            return PyErr_Format(PyExc_ValueError,
+                                "map() argument %d is longer than argument%s%d",
+                                i + 1, plural, i);
+        }
+        if (PyErr_Occurred()) {
+            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!
+    goto exit; // Cleanup
 }
 
 static PyObject *
@@ -1465,21 +1478,41 @@ map_reduce(mapobject *lz, PyObject *Py_UNUSED(ignored))
         PyTuple_SET_ITEM(args, i+1, Py_NewRef(it));
     }
 
+    if (lz->strict) {
+        Py_BuildValue("ONO", Py_TYPE(lz), args, Py_True);
+    }
     return Py_BuildValue("ON", Py_TYPE(lz), args);
 }
 
+PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
+
+static PyObject *
+map_setstate(mapobject *lz, PyObject *state)
+{
+    int strict = PyObject_IsTrue(state);
+    if (strict < 0) {
+        return NULL;
+    }
+    lz->strict = strict;
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef map_methods[] = {
     {"__reduce__", _PyCFunction_CAST(map_reduce), METH_NOARGS, reduce_doc},
+    {"__setstate__", (PyCFunction)map_setstate, METH_O, setstate_doc},
     {NULL,           NULL}           /* sentinel */
 };
 
 
 PyDoc_STRVAR(map_doc,
-"map(function, /, *iterables)\n\
+"map(function, /, *iterables, strict=False)\n\
 --\n\
 \n\
 Make an iterator that computes the function using arguments from\n\
-each of the iterables.  Stops when the shortest iterable is exhausted.");
+each of the iterables.  Stops when the shortest iterable is exhausted.\n\
+\n\
+If strict is true and one of the arguments is exhausted before the others,\n\
+raise a ValueError.");
 
 PyTypeObject PyMap_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
@@ -1523,7 +1556,6 @@ PyTypeObject PyMap_Type = {
     PyType_GenericAlloc,                /* tp_alloc */
     map_new,                            /* tp_new */
     PyObject_GC_Del,                    /* tp_free */
-    .tp_vectorcall = (vectorcallfunc)map_vectorcall
 };
 
 
@@ -2974,8 +3006,6 @@ zip_reduce(zipobject *lz, PyObject *Py_UNUSED(ignored))
     return PyTuple_Pack(2, Py_TYPE(lz), lz->ittuple);
 }
 
-PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
-
 static PyObject *
 zip_setstate(zipobject *lz, PyObject *state)
 {