Index: Doc/library/itertools.rst =================================================================== --- Doc/library/itertools.rst (revision 59544) +++ Doc/library/itertools.rst (working copy) @@ -125,6 +125,36 @@ yield x +.. function:: group(iterable[, n][, step]) + + Make an iterator that returns *n* consecutive elements from the *iterable* at a + time. The *step* is how far the start of each group progresses from the start of + the previous group. If not specified, *n* defaults to 2, and *step* defaults to + the value of *n*. If the end of the *iterable* is reached before completing a + full group of *n* items, the incomplete group will be dropped. Equivalent to:: + + def group(iterable, n=2, step=None): + if n <= 0: + return + if not step or step < 0: + step = n + head = 0 + buffer = [None] * n + i = iter(collection) + for j in range(n): + buffer[j] = i.next() + yield tuple(buffer) + step_cnt = 0 + while True: + buffer[head] = i.next() + head = (head+1)%n + step_cnt += 1 + if step_cnt >= step: + step_cnt = 0 + yield tuple(buffer[head:]) + tuple(buffer[:head]) + + + .. function:: groupby(iterable[, key]) Make an iterator that returns consecutive keys and groups from the *iterable*. Index: Lib/test/test_itertools.py =================================================================== --- Lib/test/test_itertools.py (revision 59544) +++ Lib/test/test_itertools.py (working copy) @@ -79,6 +79,21 @@ self.assertRaises(TypeError, cycle, 5) self.assertEqual(list(islice(cycle(gen3()),10)), [0,1,2,0,1,2,0,1,2,0]) + def test_group(self): + self.assertEqual(list(group([])), []) + self.assertEqual(list(group([], n=3)), []) + self.assertEqual(list(group([], step=3)), []) + self.assertRaises(TypeError, group, None) + self.assertRaises(TypeError, group, [], 'abc') + self.assertRaises(TypeError, group, [], 2, 'abc') + + s = [1, 2, 3, 4, 5, 6] + self.assertEqual(list(group(s)), [(1, 2), (3, 4), (5, 6)]) + self.assertEqual(list(group(s)), [(1, 2), (3, 4), (5, 6)]) + self.assertEqual(list(group(s, n=3)), [(1, 2, 3), (4, 5, 6)]) + self.assertEqual(list(group(s, step=3)), [(1, 2), (4, 5)]) + self.assertEqual(list(group(s, n=3, step=2)), [(1, 2, 3), (3, 4, 5)]) + def test_groupby(self): # Check whether it accepts arguments correctly self.assertEqual([], list(groupby([]))) @@ -498,6 +513,10 @@ a = [] self.makecycle(dropwhile(bool, [0, a, a]), a) + def test_group(self): + a = [] + self.makecycle(group([a]*2), a) + def test_groupby(self): a = [] self.makecycle(groupby([a]*2, lambda x:x), a) @@ -633,6 +652,14 @@ self.assertRaises(TypeError, list, cycle(N(s))) self.assertRaises(ZeroDivisionError, list, cycle(E(s))) + def test_group(self): + for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): + for g in (G, I, Ig, S, L, R): + self.assertEqual(list(group(g(s))), zip(*(iter(g(s)),)*2)) + self.assertRaises(TypeError, group, X(s)) + self.assertRaises(TypeError, list, group(N(s))) + self.assertRaises(ZeroDivisionError, list, group(E(s))) + def test_groupby(self): for s in (range(10), range(0), range(1000), (7,11), xrange(2000,2200,5)): for g in (G, I, Ig, S, L, R): Index: Modules/itertoolsmodule.c =================================================================== --- Modules/itertoolsmodule.c (revision 59544) +++ Modules/itertoolsmodule.c (working copy) @@ -2745,6 +2745,170 @@ PyObject_GC_Del, /* tp_free */ }; +/* group object *************************************************************/ + +typedef struct { + PyObject_HEAD + PyObject *it; + PyObject *buffer; + Py_ssize_t n; + Py_ssize_t step; + Py_ssize_t head; +} groupobject; + +static PyTypeObject group_type; + +static PyObject * +group_new(PyTypeObject *type, PyObject *args, PyObject *kwds) +{ + static char *kwargs[] = {"iterable", "n", "step", NULL}; + groupobject *obj; + Py_ssize_t n = 2; + Py_ssize_t step = 0; + PyObject *collection, *it; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|nn:group", kwargs, + &collection, &n, &step)) + return NULL; + + it = PyObject_GetIter(collection); + if (it == NULL) + return NULL; + if (step <= 0) + step = n; + + obj = (groupobject *)type->tp_alloc(type, 0); + if (obj == NULL) { + Py_DECREF(it); + return NULL; + } + obj->it = it; + obj->buffer = NULL; + obj->n = n; + obj->step = step; + obj->head = 0; + + return (PyObject *)obj; +} + +static void +group_dealloc(groupobject *self) +{ + PyObject_GC_UnTrack(self); + Py_XDECREF(self->it); + Py_XDECREF(self->buffer); + Py_Type(self)->tp_free(self); +} + +static int +group_traverse(groupobject *self, visitproc visit, void *arg) +{ + Py_VISIT(self->it); + Py_VISIT(self->buffer); + return 0; +} + +static PyObject * +group_next(groupobject *self) +{ + int i; + PyObject *item, *buffer, *ret; + + if (self->buffer == NULL) { + /* Create and fill the buffer the first time this is called. */ + buffer = PyList_New(self->n); + if (buffer == NULL) + return NULL; + for (i = 0; i < self->n; ++i) { + item = PyIter_Next(self->it); + if (item == NULL) { + Py_DECREF(buffer); + return NULL; + } + PyList_SET_ITEM(buffer, i, item); + } + self->buffer = buffer; + ret = PySequence_Tuple(buffer); + if (ret == NULL) { + /* The buffer will be cleaned up when the group object is + collected. + */ + return NULL; + } + return ret; + } else { + /* Now circle around the buffer 'step' times on each iteration. */ + for (i = 0; i < self->step; ++i) { + item = PyIter_Next(self->it); + if (item == NULL) + return NULL; + if (PyList_SetItem(self->buffer, self->head, item)) + return NULL; + self->head = (self->head + 1) % self->n; + } + ret = PyTuple_New(self->n); + if (ret == NULL) + return NULL; + for (i = 0; i < self->n; ++i) { + item = PyList_GET_ITEM(self->buffer, (i + self->head) % self->n); + if (item == NULL) { + Py_DECREF(ret); + return NULL; + } + PyTuple_SET_ITEM(ret, i, item); + } + return ret; + } +} + +PyDoc_STRVAR(group_doc, +"group(iterable[, n[, step]]) -> create an iterator which returns\n\ +subgroups of the collection.\n"); + +static PyTypeObject group_type = { + PyVarObject_HEAD_INIT(NULL, 0) + "itertools.group", /* tp_name */ + sizeof(groupobject), /* tp_basicsize */ + 0, /* tp_itemsize */ + /* methods */ + (destructor)group_dealloc, /* tp_dealloc */ + 0, /* tp_print */ + 0, /* tp_getattr */ + 0, /* tp_setattr */ + 0, /* tp_compare */ + 0, /* tp_repr */ + 0, /* tp_as_number */ + 0, /* tp_as_sequence */ + 0, /* tp_as_mapping */ + 0, /* tp_hash */ + 0, /* tp_call */ + 0, /* tp_str */ + PyObject_GenericGetAttr, /* tp_getattro */ + 0, /* tp_setattro */ + 0, /* tp_as_buffer */ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | + Py_TPFLAGS_BASETYPE, /* tp_flags */ + group_doc, /* tp_doc */ + (traverseproc)group_traverse, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + PyObject_SelfIter, /* tp_iter */ + (iternextfunc)group_next, /* tp_iternext */ + 0, /* tp_methods */ + 0, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + 0, /* tp_init */ + 0, /* tp_alloc */ + group_new, /* tp_new */ + PyObject_GC_Del, /* tp_free */ +}; + /* module level code ********************************************************/ PyDoc_STRVAR(module_doc, @@ -2769,6 +2933,7 @@ takewhile(pred, seq) --> seq[0], seq[1], until pred fails\n\ dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails\n\ groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)\n\ +group(iterable, n=2[, step]) --> yields elements in groups of size n\n\ "); @@ -2798,6 +2963,7 @@ &iziplongest_type, &repeat_type, &groupby_type, + &group_type, NULL };