summaryrefslogtreecommitdiffstats
path: root/Modules/itertoolsmodule.c
diff options
context:
space:
mode:
Diffstat (limited to 'Modules/itertoolsmodule.c')
-rw-r--r--Modules/itertoolsmodule.c30
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\