Index: Doc/library/itertools.rst =================================================================== --- Doc/library/itertools.rst (revision 68858) +++ Doc/library/itertools.rst (working copy) @@ -546,7 +546,27 @@ .. versionadded:: 2.4 +.. function:: fixlen(iterable, length[, value=None]) + Return an iterator fixing *iterable*\ 's length to *length* by either cutting + elements off or adding *value* to the iterable. Equivalent to:: + + def fixlen(iterable, length, value=None): + # fixlen(range(3), 2) --> 0 1 + # fixlen(range(3), 3) --> 0 1 2 + # fixlen(range(3), 4) --> 0 1 2 None + # fixlen(range(3), 4, False) --> 0 1 2 False + return islice(chain(iterable, repeat(value)), length) + + .. note:: The iterable is checked to raise `StopIteration`:exc: on each and + every element. When it decides to return elements after it has + already raised `StopIteration`:exc: once the generator will + happily return that newly arrived element. + Thus, `fixlen` works with recontinuable iterables. + + .. versionadded:: 2.7 + + .. _itertools-example: Examples Index: Lib/test/test_itertools.py =================================================================== --- Lib/test/test_itertools.py (revision 68858) +++ Lib/test/test_itertools.py (working copy) @@ -687,6 +687,19 @@ self.assertRaises(StopIteration, f(lambda x:x, []).next) self.assertRaises(StopIteration, f(lambda x:x, StopNow()).next) + def test_fixlen(self): + self.assertEqual(list(fixlen(range(3), 2)), [0, 1]) + self.assertEqual(list(fixlen(range(3), 3)), [0, 1, 2]) + self.assertEqual(list(fixlen(range(3), 4)), [0, 1, 2, None]) + self.assertEqual(list(fixlen(range(3), 4, False)), [0, 1, 2, False]) + + self.assertEqual(list(fixlen([], 0)), []) + self.assertEqual(list(fixlen([], 2)), [None, None]) + + self.assertRaises(TypeError, fixlen) + self.assertRaises(TypeError, fixlen, None) + self.assertRaises(TypeError, fixlen, (), None) + class TestExamples(unittest.TestCase): def test_chain(self): @@ -922,7 +935,19 @@ 'Test multiple tiers of iterators' return chain(imap(lambda x:x, R(Ig(G(seqn))))) +class Ir: + def __init__(self): + self.status = 0 + def __iter__(self): + return self + def next(self): + self.status += 1 + if self.status % 2: + return self.status + else: + raise StopIteration + class TestVariousIteratorArgs(unittest.TestCase): def test_chain(self): @@ -1053,6 +1078,9 @@ self.assertRaises(TypeError, list, tee(N(s))[0]) self.assertRaises(ZeroDivisionError, list, tee(E(s))[0]) + def test_fixlen(self): + self.assertEqual(list(fixlen(Ir(), 4)), [1, None, 3, None]) + class LengthTransparency(unittest.TestCase): def test_repeat(self): Index: Modules/itertoolsmodule.c =================================================================== --- Modules/itertoolsmodule.c (revision 68858) +++ Modules/itertoolsmodule.c (working copy) @@ -3505,6 +3505,138 @@ PyObject_GC_Del, /* tp_free */ }; +/* fixlen object ************************************************************/ + +typedef struct { + PyObject_HEAD + PyObject *it; + PyObject *val; + Py_ssize_t len; +} fixlenobject; + +static PyObject * +fixlen_new(PyTypeObject *type, PyObject *args, PyObject *kwds) +{ + fixlenobject *lz; + PyObject *iterable = NULL; + PyObject *value = NULL; + Py_ssize_t length; + static char *kwargs[] = {"iterable", "length", "value", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "On|O:fixlen", kwargs, + &iterable, &length, &value)) + return NULL; + + if (value == NULL) { + Py_INCREF(Py_None); + value = Py_None; + } + + lz = (fixlenobject *)type->tp_alloc(type, 0); + if (lz == NULL) { + Py_DECREF(iterable); + Py_DECREF(value); + return NULL; + } + + lz->it = PyObject_GetIter(iterable); + if (lz->it == NULL) { + Py_DECREF(lz); + return NULL; + } + lz->len = length; + lz->val = value; + + return (PyObject *)lz; +} + +static PyObject * +fixlen_next(fixlenobject *lz) +{ + PyObject *item; + + if (lz->len <= 0) + return NULL; + lz->len--; + + item = PyIter_Next(lz->it); + if (item != NULL) + return item; + else { + Py_INCREF(lz->val); + return lz->val; + } +} + +static int +fixlen_traverse(fixlenobject *lz, visitproc visit, void *arg) +{ + Py_VISIT(lz->it); + Py_VISIT(lz->val); + return 0; +} + +static void +fixlen_dealloc(fixlenobject *lz) +{ + PyObject_GC_UnTrack(lz); + Py_XDECREF(lz->it); + Py_XDECREF(lz->val); + Py_TYPE(lz)->tp_free(lz); +} + +PyDoc_STRVAR(fixlen_doc, +"fixlen(iterable, length, value=None) --> fixlen object\n\ +\n\ +Return elements from iterable until either length is reached or it is\n\ +exhausted. If iterable cannot satisfy the length condition, return the value\n\ +until length is reached."); + +static PyTypeObject fixlen_type = { + PyVarObject_HEAD_INIT(NULL, 0) + "itertools.fixlen", /* tp_name */ + sizeof(fixlenobject), /* tp_basicsize */ + 0, /* tp_itemsize */ + /* methods */ + (destructor)fixlen_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 */ + fixlen_doc, /* tp_doc */ + (traverseproc)fixlen_traverse, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + PyObject_SelfIter, /* tp_iter */ + (iternextfunc)fixlen_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 */ + fixlen_new, /* tp_new */ + PyObject_GC_Del, /* tp_free */ +}; + + /* module level code ********************************************************/ PyDoc_STRVAR(module_doc, @@ -3529,6 +3661,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\ +fixlen(iterable, length[, value]) --> iterable fixed to a certain length\n\ "); @@ -3561,6 +3694,7 @@ &product_type, &repeat_type, &groupby_type, + &fixlen_type, NULL };