python · MSeifert04 · Apr 11, 2017 · nascheme · Apr 28, 2017 · MSeifert04
diff --git 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
@@ -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 */
 };