Repository: aio-libs/multidict

Issue:
Implement __length_hint__() for iterators
It can help with optimizing calls like `[i for i in md.items()]` etc.

Comments:

Patch:
diff --git a/CHANGES/310.feature b/CHANGES/310.feature
new file mode 100644
index 0000000..9f0fb19
--- /dev/null
+++ b/CHANGES/310.feature
@@ -0,0 +1,1 @@
+Implement __length_hint__() for iterators
\ No newline at end of file
diff --git a/multidict/_multidict_iter.c b/multidict/_multidict_iter.c
index 41aeb71..30bb46b 100644
--- a/multidict/_multidict_iter.c
+++ b/multidict/_multidict_iter.c
@@ -156,6 +156,28 @@ multidict_iter_clear(MultidictIter *self)
     return 0;
 }
 
+static PyObject *
+multidict_iter_len(MultidictIter *self)
+{
+    return PyLong_FromSize_t(pair_list_len(&self->md->pairs));
+}
+
+PyDoc_STRVAR(length_hint_doc,
+             "Private method returning an estimate of len(list(it)).");
+
+static PyMethodDef multidict_iter_methods[] = {
+    {
+        "__length_hint__",
+        (PyCFunction)(void(*)(void))multidict_iter_len,
+        METH_NOARGS,
+        length_hint_doc
+    },
+    {
+        NULL,
+        NULL
+    }   /* sentinel */
+};
+
 /***********************************************************************/
 
 /* We link this module statically for convenience.  If compiled as a shared
@@ -177,6 +199,7 @@ static PyTypeObject multidict_items_iter_type = {
     .tp_clear = (inquiry)multidict_iter_clear,
     .tp_iter = PyObject_SelfIter,
     .tp_iternext = (iternextfunc)multidict_items_iter_iternext,
+    .tp_methods = multidict_iter_methods,
 };
 
 static PyTypeObject multidict_values_iter_type = {
@@ -189,6 +212,7 @@ static PyTypeObject multidict_values_iter_type = {
     .tp_clear = (inquiry)multidict_iter_clear,
     .tp_iter = PyObject_SelfIter,
     .tp_iternext = (iternextfunc)multidict_values_iter_iternext,
+    .tp_methods = multidict_iter_methods,
 };
 
 static PyTypeObject multidict_keys_iter_type = {
@@ -201,6 +225,7 @@ static PyTypeObject multidict_keys_iter_type = {
     .tp_clear = (inquiry)multidict_iter_clear,
     .tp_iter = PyObject_SelfIter,
     .tp_iternext = (iternextfunc)multidict_keys_iter_iternext,
+    .tp_methods = multidict_iter_methods,
 };
 
 int
diff --git a/multidict/_multidict_py.py b/multidict/_multidict_py.py
index 1d5736d..04e96cd 100644
--- a/multidict/_multidict_py.py
+++ b/multidict/_multidict_py.py
@@ -404,6 +404,23 @@ class CIMultiDict(MultiDict):
         return key.title()
 
 
+class _Iter:
+    __slots__ = ('_size', '_iter')
+
+    def __init__(self, size, iterator):
+        self._size = size
+        self._iter = iterator
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return next(self._iter)
+
+    def __length_hint__(self):
+        return self._size
+
+
 class _ViewBase:
     def __init__(self, impl):
         self._impl = impl
@@ -423,6 +440,9 @@ class _ItemsView(_ViewBase, abc.ItemsView):
         return False
 
     def __iter__(self):
+        return _Iter(len(self), self._iter())
+
+    def _iter(self):
         for i, k, v in self._impl._items:
             if self._version != self._impl._version:
                 raise RuntimeError("Dictionary changed during iteration")
@@ -444,6 +464,9 @@ class _ValuesView(_ViewBase, abc.ValuesView):
         return False
 
     def __iter__(self):
+        return _Iter(len(self), self._iter())
+
+    def _iter(self):
         for item in self._impl._items:
             if self._version != self._impl._version:
                 raise RuntimeError("Dictionary changed during iteration")
@@ -465,6 +488,9 @@ class _KeysView(_ViewBase, abc.KeysView):
         return False
 
     def __iter__(self):
+        return _Iter(len(self), self._iter())
+
+    def _iter(self):
         for item in self._impl._items:
             if self._version != self._impl._version:
                 raise RuntimeError("Dictionary changed during iteration")


Test patch:
diff --git a/tests/test_multidict.py b/tests/test_multidict.py
index 92dc62b..6822eca 100644
--- a/tests/test_multidict.py
+++ b/tests/test_multidict.py
@@ -320,6 +320,21 @@ class BaseMultiDictTest:
         assert called
         del wr
 
+    def test_iter_length_hint_keys(self, cls):
+        md = cls(a=1, b=2)
+        it = iter(md.keys())
+        assert it.__length_hint__() == 2
+
+    def test_iter_length_hint_items(self, cls):
+        md = cls(a=1, b=2)
+        it = iter(md.items())
+        assert it.__length_hint__() == 2
+
+    def test_iter_length_hint_values(self, cls):
+        md = cls(a=1, b=2)
+        it = iter(md.values())
+        assert it.__length_hint__() == 2
+
 
 class TestMultiDict(BaseMultiDictTest):
     @pytest.fixture(params=["MultiDict", ("MultiDict", "MultiDictProxy")])