diff options
Diffstat (limited to 'Modules/itertoolsmodule.c')
-rw-r--r-- | Modules/itertoolsmodule.c | 30 |
1 files changed, 21 insertions, 9 deletions
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 8b6fa85..15b0c17 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -626,6 +626,7 @@ tee(PyObject *self, PyObject *args) { Py_ssize_t i, n=2; PyObject *it, *iterable, *copyable, *result; + _Py_IDENTIFIER(__copy__); if (!PyArg_ParseTuple(args, "O|n", &iterable, &n)) return NULL; @@ -643,7 +644,7 @@ tee(PyObject *self, PyObject *args) Py_DECREF(result); return NULL; } - if (!PyObject_HasAttrString(it, "__copy__")) { + if (!_PyObject_HasAttrId(it, &PyId___copy__)) { copyable = tee_fromiterable(it); Py_DECREF(it); if (copyable == NULL) { @@ -654,7 +655,8 @@ tee(PyObject *self, PyObject *args) copyable = it; PyTuple_SET_ITEM(result, 0, copyable); for (i=1 ; i<n ; i++) { - copyable = PyObject_CallMethod(copyable, "__copy__", NULL); + + copyable = _PyObject_CallMethodId(copyable, &PyId___copy__, NULL); if (copyable == NULL) { Py_DECREF(result); return NULL; @@ -2592,6 +2594,7 @@ typedef struct { PyObject_HEAD PyObject *total; PyObject *it; + PyObject *binop; } accumulateobject; static PyTypeObject accumulate_type; @@ -2599,12 +2602,14 @@ static PyTypeObject accumulate_type; static PyObject * accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds) { - static char *kwargs[] = {"iterable", NULL}; + static char *kwargs[] = {"iterable", "func", NULL}; PyObject *iterable; PyObject *it; + PyObject *binop = NULL; accumulateobject *lz; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:accumulate", kwargs, &iterable)) + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:accumulate", + kwargs, &iterable, &binop)) return NULL; /* Get iterator. */ @@ -2619,6 +2624,8 @@ accumulate_new(PyTypeObject *type, PyObject *args, PyObject *kwds) return NULL; } + Py_XINCREF(binop); + lz->binop = binop; lz->total = NULL; lz->it = it; return (PyObject *)lz; @@ -2628,6 +2635,7 @@ static void accumulate_dealloc(accumulateobject *lz) { PyObject_GC_UnTrack(lz); + Py_XDECREF(lz->binop); Py_XDECREF(lz->total); Py_XDECREF(lz->it); Py_TYPE(lz)->tp_free(lz); @@ -2636,6 +2644,7 @@ accumulate_dealloc(accumulateobject *lz) static int accumulate_traverse(accumulateobject *lz, visitproc visit, void *arg) { + Py_VISIT(lz->binop); Py_VISIT(lz->it); Py_VISIT(lz->total); return 0; @@ -2655,8 +2664,11 @@ accumulate_next(accumulateobject *lz) lz->total = val; return lz->total; } - - newtotal = PyNumber_Add(lz->total, val); + + if (lz->binop == NULL) + newtotal = PyNumber_Add(lz->total, val); + else + newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL); Py_DECREF(val); if (newtotal == NULL) return NULL; @@ -2670,9 +2682,9 @@ accumulate_next(accumulateobject *lz) } PyDoc_STRVAR(accumulate_doc, -"accumulate(iterable) --> accumulate object\n\ +"accumulate(iterable[, func]) --> accumulate object\n\ \n\ -Return series of accumulated sums."); +Return series of accumulated sums (or other binary function results)."); static PyTypeObject accumulate_type = { PyVarObject_HEAD_INIT(NULL, 0) @@ -3633,7 +3645,7 @@ cycle(p) --> p0, p1, ... plast, p0, p1, ...\n\ repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n times\n\ \n\ Iterators terminating on the shortest input sequence:\n\ -accumulate(p, start=0) --> p0, p0+p1, p0+p1+p2\n\ +accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2\n\ chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ... \n\ compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...\n\ dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\ |