Mark map/accumulate iterators exhausted when the user callback raises StopIteration

Author: PrathmeshAdsod

Modules/itertoolsmodule.c

@@ -2983,6 +2983,7 @@ typedef struct {
     PyObject *it;
     PyObject *binop;
     PyObject *initial;
+    int finished;
     itertools_state *state;
 } accumulateobject;
 
@@ -3024,6 +3025,7 @@ itertools_accumulate_impl(PyTypeObject *type, PyObject *iterable,
     lz->total = NULL;
     lz->it = it;
     lz->initial = Py_XNewRef(initial);
+    lz->finished = 0;
     lz->state = find_state_by_type(type);
     return (PyObject *)lz;
 }
@@ -3060,6 +3062,10 @@ accumulate_next(PyObject *op)
     accumulateobject *lz = accumulateobject_CAST(op);
     PyObject *val, *newtotal;
 
+    if (lz->finished) {
+        return NULL;
+    }
+
     if (lz->initial != Py_None) {
         lz->total = lz->initial;
         lz->initial = Py_NewRef(Py_None);
@@ -3079,8 +3085,12 @@ accumulate_next(PyObject *op)
     else
         newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL);
     Py_DECREF(val);
-    if (newtotal == NULL)
+    if (newtotal == NULL) {
+        if (lz->binop != NULL && PyErr_Occurred() && PyErr_ExceptionMatches(PyExc_StopIteration)) {
+            lz->finished = 1;
+        }
         return NULL;
+    }
 
     Py_INCREF(newtotal);
     Py_SETREF(lz->total, newtotal);

Python/bltinmodule.c

@@ -1354,6 +1354,7 @@ typedef struct {
     PyObject *iters;
     PyObject *func;
     int strict;
+    int finished;
 } mapobject;
 
 #define _mapobject_CAST(op)     ((mapobject *)(op))
@@ -1411,6 +1412,7 @@ map_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     func = PyTuple_GET_ITEM(args, 0);
     lz->func = Py_NewRef(func);
     lz->strict = strict;
+    lz->finished = 0;
 
     return (PyObject *)lz;
 }
@@ -1456,6 +1458,7 @@ map_vectorcall(PyObject *type, PyObject * const*args,
     lz->iters = iters;
     lz->func = Py_NewRef(args[0]);
     lz->strict = 0;
+    lz->finished = 0;
 
     return (PyObject *)lz;
 }
@@ -1489,6 +1492,10 @@ map_next(PyObject *self)
     PyObject *result = NULL;
     PyThreadState *tstate = _PyThreadState_GET();
 
+    if (lz->finished) {
+        return NULL;
+    }
+
     const Py_ssize_t niters = PyTuple_GET_SIZE(lz->iters);
     if (niters <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) {
         stack = small_stack;
@@ -1516,6 +1523,11 @@ map_next(PyObject *self)
     }
 
     result = _PyObject_VectorcallTstate(tstate, lz->func, stack, nargs, NULL);
+    if (result == NULL && PyErr_Occurred()) {
+        if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
+            lz->finished = 1;
+        }
+    }
 
 exit:
     for (i=0; i < nargs; i++) {