diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py index d2420e9dda1e07..36a9b4d3353b96 100644 --- a/Lib/test/test_builtin.py +++ b/Lib/test/test_builtin.py @@ -17,7 +17,7 @@ import types import unittest import warnings -from operator import neg +from operator import neg, length_hint from test.support import TESTFN, unlink, check_warnings from test.support.script_helper import assert_python_ok try: @@ -837,6 +837,23 @@ def badfunc(x): raise RuntimeError self.assertRaises(RuntimeError, list, map(badfunc, range(5))) + def test_map_length_hint(self): + self.assertEqual(4, length_hint(map(int, [1]*4))) + self.assertEqual(1, length_hint(map(max, [1]*100, [1]))) + # Something without a length_hint: A generator + self.assertEqual(0, length_hint(map(max, (i for i in [1])))) + # Something that raises an error inside length_hint + class BadLengthHintIterator(object): + def __iter__(self): + return self + def __length_hint__(self): + raise ValueError + self.assertRaises(ValueError, length_hint, map(int, BadLengthHintIterator())) + + mapit = map(max, [1]*10, [1]*5) + next(mapit) + self.assertEqual(4, length_hint(mapit)) + def test_map_pickle(self): for proto in range(pickle.HIGHEST_PROTOCOL + 1): m1 = map(map_char, "Is this the real life?") @@ -1367,6 +1384,26 @@ def __getitem__(self, i): return i self.assertRaises(ValueError, list, zip(BadSeq(), BadSeq())) + def test_zip_length_hint(self): + self.assertEqual(0, length_hint(zip())) + self.assertEqual(4, length_hint(zip([1]*4))) + self.assertEqual(1, length_hint(zip([1]*100, [1]))) + # Something without a length_hint: A generator + self.assertEqual(0, length_hint(zip((i for i in [1]), [1]*10))) + + # Something that raises an error inside length_hint + class BadLengthHintIterator(object): + def __iter__(self): + return self + def __length_hint__(self): + raise ValueError + self.assertRaises(ValueError, length_hint, zip(BadLengthHintIterator())) + + zipit = zip([1]*5, [1]*10) + next(zipit) + self.assertEqual(4, length_hint(zipit)) + + def test_zip_pickle(self): a = (1, 2, 3) b = (4, 5, 6) diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c index 8ae2303e98d333..ac71d1ff94752e 100644 --- a/Python/bltinmodule.c +++ b/Python/bltinmodule.c @@ -1221,6 +1221,26 @@ map_next(mapobject *lz) return result; } +static PyObject * +map_len(mapobject *lz) +{ + Py_ssize_t length_hint = PY_SSIZE_T_MAX; + Py_ssize_t i; + + for (i = 0; length_hint > 0 && i < PyTuple_GET_SIZE(lz->iters); i++) { + PyObject *it = PyTuple_GET_ITEM(lz->iters, i); + Py_ssize_t it_len = PyObject_LengthHint(it, 0); + + if (it_len < length_hint) + length_hint = it_len; + } + + if (length_hint == -1) + return NULL; + + return PyLong_FromSsize_t(length_hint); +} + static PyObject * map_reduce(mapobject *lz) { @@ -1240,8 +1260,11 @@ map_reduce(mapobject *lz) return Py_BuildValue("ON", Py_TYPE(lz), args); } +PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it))."); + static PyMethodDef map_methods[] = { - {"__reduce__", (PyCFunction)map_reduce, METH_NOARGS, reduce_doc}, + {"__length_hint__", (PyCFunction)map_len, METH_NOARGS, length_hint_doc}, + {"__reduce__", (PyCFunction)map_reduce, METH_NOARGS, reduce_doc}, {NULL, NULL} /* sentinel */ }; @@ -2550,6 +2573,32 @@ zip_next(zipobject *lz) return result; } +static PyObject * +zip_len(zipobject *lz) +{ + Py_ssize_t tuplesize = lz->tuplesize; + Py_ssize_t length_hint = PY_SSIZE_T_MAX; + Py_ssize_t i; + + if (tuplesize == 0) { + Py_INCREF(_PyLong_Zero); + return _PyLong_Zero; + } + + for (i = 0; length_hint > 0 && i < PyTuple_GET_SIZE(lz->ittuple); i++) { + PyObject *it = PyTuple_GET_ITEM(lz->ittuple, i); + Py_ssize_t it_len = PyObject_LengthHint(it, 0); + + if (it_len < length_hint) + length_hint = it_len; + } + + if (length_hint == -1) + return NULL; + + return PyLong_FromSsize_t(length_hint); +} + static PyObject * zip_reduce(zipobject *lz) { @@ -2558,7 +2607,8 @@ zip_reduce(zipobject *lz) } static PyMethodDef zip_methods[] = { - {"__reduce__", (PyCFunction)zip_reduce, METH_NOARGS, reduce_doc}, + {"__length_hint__", (PyCFunction)zip_len, METH_NOARGS, length_hint_doc}, + {"__reduce__", (PyCFunction)zip_reduce, METH_NOARGS, reduce_doc}, {NULL, NULL} /* sentinel */ };