Safely iterate over non-dict mappings. See #127

cytoolz/dicttoolz.pxd

@@ -1,6 +1,10 @@
 from cpython.ref cimport PyObject
 
 # utility functions to perform iteration over dicts or generic mapping
+cdef class _iter_mapping:
+    cdef object it
+    cdef object cur
+
 ctypedef int (*f_map_next)(object p, Py_ssize_t *ppos, PyObject* *pkey, PyObject* *pval) except -1
 
 cdef f_map_next get_map_iter(object d, PyObject* *ptr) except NULL

cytoolz/dicttoolz.pyx

@@ -16,6 +16,20 @@ __all__ = ['merge', 'merge_with', 'valmap', 'keymap', 'itemmap', 'valfilter',
            'update_in']
 
 
+cdef class _iter_mapping:
+    """ Keep a handle on the current item to prevent memory clean up too early"""
+    def __cinit__(self, object it):
+        self.it = it
+        self.cur = None
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        self.cur = next(self.it)
+        return self.cur
+
+
 cdef int PyMapping_Next(object p, Py_ssize_t *ppos, PyObject* *pkey, PyObject* *pval) except -1:
     """Mimic "PyDict_Next" interface, but for any mapping"""
     cdef PyObject *obj
@@ -24,7 +38,7 @@ cdef int PyMapping_Next(object p, Py_ssize_t *ppos, PyObject* *pkey, PyObject* *
         return 0
     pkey[0] = <PyObject*>(<object>obj)[0]
     pval[0] = <PyObject*>(<object>obj)[1]
-    Py_XDECREF(obj)
+    Py_XDECREF(obj)  # removing this results in memory leak
     return 1
 
 
@@ -53,10 +67,10 @@ cdef f_map_next get_map_iter(object d, PyObject* *ptr) except NULL:
         val = d
         rv = &PyDict_Next_Compat
     elif hasattr(d, 'iteritems'):
-        val = iter(d.iteritems())
+        val = _iter_mapping(iter(d.iteritems()))
         rv = &PyMapping_Next
     else:
-        val = iter(d.items())
+        val = _iter_mapping(iter(d.items()))
         rv = &PyMapping_Next
     Py_INCREF(val)
     ptr[0] = <PyObject*>val

cytoolz/tests/test_dicttoolz.py

@@ -1,7 +1,9 @@
 from collections import defaultdict as _defaultdict
+import os
 from cytoolz.dicttoolz import (merge, merge_with, valmap, keymap, update_in,
                              assoc, dissoc, keyfilter, valfilter, itemmap,
                              itemfilter, assoc_in)
+from cytoolz.functoolz import identity
 from cytoolz.utils import raises
 from cytoolz.compatibility import PY3
 
@@ -250,3 +252,10 @@ class TestCustomMapping(TestDict):
     """
     D = CustomMapping
     kw = {'factory': lambda: CustomMapping()}
+
+
+def test_environ():
+    # See: https://github.com/pytoolz/cytoolz/issues/127
+    assert keymap(identity, os.environ) == os.environ
+    assert valmap(identity, os.environ) == os.environ
+    assert itemmap(identity, os.environ) == os.environ