Skip to content

Commit 5f94bd0

Browse files
gh-145300: Add __length_hint__ for itertools.islice
1 parent 72eca2a commit 5f94bd0

File tree

3 files changed

+46
-0
lines changed

3 files changed

+46
-0
lines changed

Lib/test/test_itertools.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1213,6 +1213,16 @@ def __index__(self):
12131213
self.assertEqual(list(islice(range(100), IntLike(10), IntLike(50), IntLike(5))),
12141214
list(range(10,50,5)))
12151215

1216+
# Test __length_hint__
1217+
self.assertEqual(islice(range(10), 1).__length_hint__(), 1)
1218+
self.assertEqual(islice(range(10), 1, 2).__length_hint__(), 1)
1219+
self.assertEqual(islice(range(10), 1, 2, 2).__length_hint__(), 1)
1220+
self.assertEqual(islice(range(10), 0, 4, 2).__length_hint__(), 2)
1221+
self.assertEqual(islice(range(10), 2, None, 1).__length_hint__(), 8)
1222+
self.assertEqual(islice(range(10), 2, None, 2).__length_hint__(), 4)
1223+
self.assertEqual(islice(range(10), 3, None, 2).__length_hint__(), 4)
1224+
self.assertEqual(islice(range(10), 2, 2, 1).__length_hint__(), 0)
1225+
12161226
def test_takewhile(self):
12171227
data = [1, 3, 5, 20, 2, 4, 6, 8]
12181228
self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add :meth:`~itertools.islice.__length_hint__` to :class:`itertools.islice`.

Modules/itertoolsmodule.c

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1660,6 +1660,35 @@ islice_next(PyObject *op)
16601660
return NULL;
16611661
}
16621662

1663+
static PyObject *
1664+
islice_length_hint(PyObject *op, PyObject *Py_UNUSED(dummy))
1665+
{
1666+
isliceobject *lz = isliceobject_CAST(op);
1667+
1668+
if (lz->stop >= 0 && lz->stop <= lz->next) {
1669+
return PyLong_FromSsize_t(0);
1670+
}
1671+
1672+
Py_ssize_t remaining;
1673+
if (lz->stop == -1) {
1674+
Py_ssize_t hint = PyObject_LengthHint(lz->it, 0);
1675+
if (hint < 0) {
1676+
/* propagate exception */
1677+
return NULL;
1678+
}
1679+
remaining = hint - lz->next;
1680+
} else {
1681+
remaining = lz->stop - lz->next;
1682+
}
1683+
1684+
if (remaining <= 0) {
1685+
return PyLong_FromSsize_t(0);
1686+
}
1687+
1688+
Py_ssize_t steps = (remaining + lz->step - 1) / lz->step;
1689+
return PyLong_FromSsize_t(steps);
1690+
}
1691+
16631692
PyDoc_STRVAR(islice_doc,
16641693
"islice(iterable, stop) --> islice object\n\
16651694
islice(iterable, start, stop[, step]) --> islice object\n\
@@ -1671,6 +1700,11 @@ specified as another value, step determines how many values are\n\
16711700
skipped between successive calls. Works like a slice() on a list\n\
16721701
but returns an iterator.");
16731702

1703+
static PyMethodDef islice_methods[] = {
1704+
{"__length_hint__", islice_length_hint, METH_NOARGS, NULL},
1705+
{NULL, NULL},
1706+
};
1707+
16741708
static PyType_Slot islice_slots[] = {
16751709
{Py_tp_dealloc, islice_dealloc},
16761710
{Py_tp_getattro, PyObject_GenericGetAttr},
@@ -1680,6 +1714,7 @@ static PyType_Slot islice_slots[] = {
16801714
{Py_tp_iternext, islice_next},
16811715
{Py_tp_new, islice_new},
16821716
{Py_tp_free, PyObject_GC_Del},
1717+
{Py_tp_methods, islice_methods},
16831718
{0, NULL},
16841719
};
16851720

0 commit comments

Comments
 (0)