diff --git a/cytoolz/dicttoolz.pxd b/cytoolz/dicttoolz.pxd index 4b80e9b..22d91bc 100644 --- a/cytoolz/dicttoolz.pxd +++ b/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 diff --git a/cytoolz/dicttoolz.pyx b/cytoolz/dicttoolz.pyx index aa0c213..8235899 100644 --- a/cytoolz/dicttoolz.pyx +++ b/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] = (obj)[0] pval[0] = (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] = val diff --git a/cytoolz/tests/test_dicttoolz.py b/cytoolz/tests/test_dicttoolz.py index e4de4df..78ca675 100644 --- a/cytoolz/tests/test_dicttoolz.py +++ b/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