From ae0480032d62a0c319cc4f6860ad3e171877d8e4 Mon Sep 17 00:00:00 2001
From: Lars Buitinck <larsmans@gmail.com>
Date: Thu, 1 May 2014 18:20:52 +0200
Subject: [PATCH] optimize dictionary iteration

... where it makes sense, so not in c_merge_with where there was
no speedup, only code complication. Example, before:

>>> iseven = lambda x: x % 2 == 0
>>> d = dict(enumerate(range(1, 25000)))
>>> %timeit keyfilter(iseven, d)
100 loops, best of 3: 3.19 ms per loop

After:

>>> %timeit keyfilter(iseven, d)
100 loops, best of 3: 2.85 ms per loop
---
 cytoolz/dicttoolz.pyx           | 67 +++++++++++++++++++++++++--------
 cytoolz/tests/test_none_safe.py | 12 ++----
 2 files changed, 56 insertions(+), 23 deletions(-)

diff --git a/cytoolz/dicttoolz.pyx b/cytoolz/dicttoolz.pyx
index a807fe7..7c69e93 100644
--- a/cytoolz/dicttoolz.pyx
+++ b/cytoolz/dicttoolz.pyx
@@ -1,6 +1,7 @@
 #cython: embedsignature=True
 from cpython.dict cimport (PyDict_Check, PyDict_GetItem, PyDict_Merge,
-                           PyDict_New, PyDict_SetItem, PyDict_Update)
+                           PyDict_New, PyDict_Next, PyDict_SetItem,
+                           PyDict_Update)
 from cpython.exc cimport PyErr_Clear, PyErr_GivenExceptionMatches, PyErr_Occurred
 from cpython.list cimport PyList_Append, PyList_New
 from cpython.ref cimport PyObject
@@ -94,10 +95,19 @@ cpdef dict valmap(object func, dict d):
     See Also:
         keymap
     """
-    cdef dict rv
+    cdef:
+        dict rv
+        Py_ssize_t pos
+        PyObject *k
+        PyObject *v
+
+    if d is None:
+        raise TypeError("expected dict, got None")
+
     rv = PyDict_New()
-    for k, v in d.iteritems():
-       PyDict_SetItem(rv, k, func(v))
+    pos = 0
+    while PyDict_Next(d, &pos, &k, &v):
+       PyDict_SetItem(rv, <object>k, func(<object>v))
     return rv
 
 
@@ -112,10 +122,19 @@ cpdef dict keymap(object func, dict d):
     See Also:
         valmap
     """
-    cdef dict rv
+    cdef:
+        dict rv
+        Py_ssize_t pos
+        PyObject *k
+        PyObject *v
+
+    if d is None:
+        raise TypeError("expected dict, got None")
+
     rv = PyDict_New()
-    for k, v in d.iteritems():
-       PyDict_SetItem(rv, func(k), v)
+    pos = 0
+    while PyDict_Next(d, &pos, &k, &v):
+       PyDict_SetItem(rv, func(<object>k), <object>v)
     return rv
 
 
@@ -132,11 +151,20 @@ cpdef dict valfilter(object predicate, dict d):
         keyfilter
         valmap
     """
-    cdef dict rv
+    cdef:
+        dict rv
+        Py_ssize_t pos
+        PyObject *k
+        PyObject *v
+
+    if d is None:
+        raise TypeError("expected dict, got None")
+
     rv = PyDict_New()
-    for k, v in d.iteritems():
-        if predicate(v):
-            PyDict_SetItem(rv, k, v)
+    pos = 0
+    while PyDict_Next(d, &pos, &k, &v):
+        if predicate(<object>v):
+            PyDict_SetItem(rv, <object>k, <object>v)
     return rv
 
 
@@ -153,11 +181,20 @@ cpdef dict keyfilter(object predicate, dict d):
         valfilter
         keymap
     """
-    cdef dict rv
+    cdef:
+        dict rv
+        Py_ssize_t pos
+        PyObject *k
+        PyObject *v
+
+    if d is None:
+        raise TypeError("expected dict, got None")
+
     rv = PyDict_New()
-    for k, v in d.iteritems():
-        if predicate(k):
-            PyDict_SetItem(rv, k, v)
+    pos = 0
+    while PyDict_Next(d, &pos, &k, &v):
+        if predicate(<object>k):
+            PyDict_SetItem(rv, <object>k, <object>v)
     return rv
 
 
diff --git a/cytoolz/tests/test_none_safe.py b/cytoolz/tests/test_none_safe.py
index 4edfeb1..b65175d 100644
--- a/cytoolz/tests/test_none_safe.py
+++ b/cytoolz/tests/test_none_safe.py
@@ -53,15 +53,13 @@ def test_dicttoolz():
     tested.append('get_in')
 
     assert raises(TypeError, lambda: keyfilter(None, {1: 2}))
-    assert raises((TypeError, AttributeError),
-                  lambda: keyfilter(identity, None))
+    assert raises(TypeError, lambda: keyfilter(identity, None))
     tested.append('keyfilter')
 
     # XXX
     assert (raises(TypeError, lambda: keymap(None, {1: 2})) or
             keymap(None, {1: 2}) == {(1,): 2})
-    assert raises((TypeError, AttributeError),
-                  lambda: keymap(identity, None))
+    assert raises(TypeError, lambda: keymap(identity, None))
     tested.append('keymap')
 
     assert raises(TypeError, lambda: merge(None))
@@ -81,15 +79,13 @@ def test_dicttoolz():
     tested.append('update_in')
 
     assert raises(TypeError, lambda: valfilter(None, {1: 2}))
-    assert raises((TypeError, AttributeError),
-                  lambda: valfilter(identity, None))
+    assert raises(TypeError, lambda: valfilter(identity, None))
     tested.append('valfilter')
 
     # XXX
     assert (raises(TypeError, lambda: valmap(None, {1: 2})) or
             valmap(None, {1: 2}) == {1: (2,)})
-    assert raises((TypeError, AttributeError),
-                  lambda: valmap(identity, None))
+    assert raises(TypeError, lambda: valmap(identity, None))
     tested.append('valmap')
 
     s1 = set(tested)