Index: Doc/lib/libitertools.tex =================================================================== --- Doc/lib/libitertools.tex (Revision 56175) +++ Doc/lib/libitertools.tex (Arbeitskopie) @@ -243,6 +243,17 @@ \end{verbatim} \end{funcdesc} +\begin{funcdesc}{getitem}{iterable, index, \optional{default}} + Returns the \var{index}'th element from the iterable. \var{index} may be + negative to count from the end. E.g. 0 returns the first element produced by + the iterator, 1 the second, -1 the last one etc. If \var{index} is negative + the iterator will be completely exhausted, if it's positive it will be + exhausted up to the \var{index}'th element. If the iterator doesn't produce + that many elements \exception{IndexError} will be raised, except when + \var{default} is given, in which case \var{default} will be returned. + \versionadded{2.6} +\end{funcdesc} + \begin{funcdesc}{islice}{iterable, \optional{start,} stop \optional{, step}} Make an iterator that returns selected elements from the iterable. If \var{start} is non-zero, then elements from the iterable are skipped Index: Lib/test/test_itertools.py =================================================================== --- Lib/test/test_itertools.py (Revision 56175) +++ Lib/test/test_itertools.py (Arbeitskopie) @@ -602,6 +602,11 @@ 'Test multiple tiers of iterators' return chain(imap(lambda x:x, R(Ig(G(seqn))))) +def Err(n): + 'generator producing an error after n items' + for i in xrange(n): + yield i + raise SyntaxError class TestVariousIteratorArgs(unittest.TestCase): @@ -727,6 +732,30 @@ self.assertRaises(TypeError, list, tee(N(s))[0]) self.assertRaises(ZeroDivisionError, list, tee(E(s))[0]) + def test_getitem(self): + iterable = [17, 23, 37] + + # Wrong arguments + self.assertRaises(TypeError, getitem) + self.assertRaises(TypeError, getitem, []) + self.assertRaises(TypeError, getitem, 42, 42) + + # Non-negative index + self.assertEqual(getitem(iterable, 0), 17) + self.assertEqual(getitem(iterable, 2), 37) + self.assertRaises(IndexError, getitem, iterable, 3) + self.assertEqual(getitem(iterable, 3, 42), 42) + self.assertEqual(getitem(Err(10), 9), 9) + self.assertRaises(SyntaxError, getitem, Err(10), 10) + + # Negative index + self.assertEqual(getitem(iterable, -1), 37) + self.assertEqual(getitem(iterable, -3), 17) + self.assertRaises(IndexError, getitem, iterable, -4) + self.assertEqual(getitem(iterable, -4, 42), 42) + # iterator is always exhausted + self.assertRaises(SyntaxError, getitem, Err(10), -1) + class LengthTransparency(unittest.TestCase): def test_repeat(self): Index: Modules/itertoolsmodule.c =================================================================== --- Modules/itertoolsmodule.c (Revision 56175) +++ Modules/itertoolsmodule.c (Arbeitskopie) @@ -1100,6 +1100,117 @@ /* islice object ************************************************************/ +static PyObject * +getitem(PyObject *self, PyObject *args) +{ + PyObject *iterable; + Py_ssize_t index; + PyObject *defaultobj = NULL; + PyObject *iter; + if (!PyArg_ParseTuple(args, "On|O:getitem", &iterable, &index, &defaultobj)) + return NULL; + + iter = PyObject_GetIter(iterable); + if (!iter) + return NULL; + if (index >= 0) { + for (;;) { + PyObject *element = PyIter_Next(iter); + if (!element){ + Py_DECREF(iter); + if (PyErr_Occurred()) + return NULL; + if (defaultobj) { + PyErr_Clear(); + Py_INCREF(defaultobj); + } + else + PyErr_SetString(PyExc_IndexError, + "iterator didn't produce enough elements"); + return defaultobj; + } + if (!index){ + /* this is the index'th element => return it */ + Py_DECREF(iter); + return element; + } + Py_DECREF(element); + --index; + } + return NULL; + } else { + Py_ssize_t size; + Py_ssize_t count; /* used by the deallocation code */ + PyObject *result = NULL; + /* We keep the last abs(index) elements in this circular buffer */ + PyObject **buffer; + Py_ssize_t headindex = 0; + index = -index; + size = index*sizeof(PyObject *); + if (size/sizeof(PyObject *) != index) { + PyErr_SetString(PyExc_OverflowError, "index too large"); + return NULL; + } + buffer = (PyObject **)PyMem_Malloc(size); + if (!buffer) + return NULL; + memset(buffer, 0, size); + for (;;) { + PyObject *element = PyIter_Next(iter); + if (++headindex >= index) + headindex = 0; + if (!element){ + Py_DECREF(iter); + if (PyErr_Occurred()) + goto finished; + /* iterator exhausted, so check if there's an object in the *next* + * slot in the ringbuffer (which is the "oldest" object we fetched + * from the iterator) */ + if (buffer[headindex]) { + result = buffer[headindex]; + /* because the dealloc code will do a decref for all objects */ + Py_INCREF(result); + goto finished; + } + if (defaultobj) { + PyErr_Clear(); + Py_INCREF(defaultobj); + result = defaultobj; + } + else + PyErr_SetString(PyExc_IndexError, + "iterator didn't produce enough elements"); + goto finished; + } else { + /* drop oldest object from the ringbuffer */ + if (buffer[headindex]) { + Py_DECREF(buffer[headindex]); + } + buffer[headindex] = element; + } + } + finished: + for (count = index; count; --count) { + if (--headindex < 0) + headindex = index-1; + if (!buffer[headindex]) + break; + Py_DECREF(buffer[headindex]); + } + PyMem_Free(buffer); + return result; + } +} + +PyDoc_STRVAR(getitem_doc, +"getitem(iterable, index [, default]) --> object\n\ +\n\ +Return the index'th element from iterable. If index is negative it is relative\n\ +to the end. If the iterator doesn't produce that many elements and default is\n\ +given it will be returned. Otherwise an IndexError will be raised."); + +/* islice object ************************************************************/ + typedef struct { PyObject_HEAD PyObject *it; @@ -2725,11 +2836,14 @@ 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\ +Functions consuming iterators:\n\ +getitem(iterable, index[, default]) --> object\n\ "); static PyMethodDef module_methods[] = { {"tee", (PyCFunction)tee, METH_VARARGS, tee_doc}, + {"getitem", (PyCFunction)getitem, METH_VARARGS, getitem_doc}, {NULL, NULL} /* sentinel */ };