[pypy-svn] r74571 - in pypy/trunk/pypy/module/cpyext: . test

[afa at codespeak.net]

Author: afa
Date: Wed May 19 12:07:12 2010
New Revision: 74571

Modified:
   pypy/trunk/pypy/module/cpyext/dictobject.py
   pypy/trunk/pypy/module/cpyext/test/test_dictobject.py
Log:
Implement (very inefficiently) PyDict_Next


Modified: pypy/trunk/pypy/module/cpyext/dictobject.py
==============================================================================
--- pypy/trunk/pypy/module/cpyext/dictobject.py	(original)
+++ pypy/trunk/pypy/module/cpyext/dictobject.py	Wed May 19 12:07:12 2010
@@ -3,6 +3,7 @@
     cpython_api, CANNOT_FAIL, build_type_checkers, Py_ssize_t,
     Py_ssize_tP, CONST_STRING)
 from pypy.module.cpyext.pyobject import PyObject, PyObjectP, borrow_from
+from pypy.module.cpyext.pyobject import RefcountState
 from pypy.module.cpyext.pyerrors import PyErr_BadInternalCall
 from pypy.interpreter.error import OperationError
 
@@ -104,7 +105,7 @@
     return space.call_method(w_obj, "items")
 
 @cpython_api([PyObject, Py_ssize_tP, PyObjectP, PyObjectP], rffi.INT_real, error=CANNOT_FAIL)
-def PyDict_Next(space, w_obj, ppos, pkey, pvalue):
+def PyDict_Next(space, w_dict, ppos, pkey, pvalue):
     """Iterate over all key-value pairs in the dictionary p.  The
     Py_ssize_t referred to by ppos must be initialized to 0
     prior to the first call to this function to start the iteration; the
@@ -145,8 +146,30 @@
         }
         Py_DECREF(o);
     }"""
-    if w_obj is None:
+    if w_dict is None:
         return 0
-    raise NotImplementedError
+
+    # Note: this is not efficient. Storing an iterator would probably
+    # work, but we can't work out how to not leak it if iteration does
+    # not complete.
+
+    try:
+        w_iter = space.call_method(w_dict, "iteritems")
+        pos = ppos[0]
+        while pos:
+            space.call_method(w_iter, "next")
+            pos -= 1
+
+        w_item = space.call_method(w_iter, "next")
+        w_key, w_value = space.fixedview(w_item, 2)
+        state = space.fromcache(RefcountState)
+        pkey[0]   = state.make_borrowed(w_dict, w_key)
+        pvalue[0] = state.make_borrowed(w_dict, w_value)
+        ppos[0] += 1
+    except OperationError, e:
+        if not e.match(space, space.w_StopIteration):
+            raise
+        return 0
+    return 1
 
 

Modified: pypy/trunk/pypy/module/cpyext/test/test_dictobject.py
==============================================================================
--- pypy/trunk/pypy/module/cpyext/test/test_dictobject.py	(original)
+++ pypy/trunk/pypy/module/cpyext/test/test_dictobject.py	Wed May 19 12:07:12 2010
@@ -1,5 +1,7 @@
 from pypy.rpython.lltypesystem import rffi, lltype
 from pypy.module.cpyext.test.test_api import BaseApiTest
+from pypy.module.cpyext.api import Py_ssize_tP, PyObjectP
+from pypy.module.cpyext.pyobject import make_ref, from_ref
 
 class TestDictObject(BaseApiTest):
     def test_dict(self, space, api):
@@ -71,3 +73,28 @@
 
         api.PyDict_Update(w_d, w_d2)
         assert space.unwrap(w_d) == dict(a='b', c='d', e='f')
+
+    def test_iter(self, space, api):
+        w_dict = space.sys.getdict()
+        py_dict = make_ref(space, w_dict)
+
+        ppos = lltype.malloc(Py_ssize_tP.TO, 1, flavor='raw')
+        ppos[0] = 0
+        pkey = lltype.malloc(PyObjectP.TO, 1, flavor='raw')
+        pvalue = lltype.malloc(PyObjectP.TO, 1, flavor='raw')
+
+        try:
+            w_copy = space.newdict()
+            while api.PyDict_Next(w_dict, ppos, pkey, pvalue):
+                w_key = from_ref(space, pkey[0])
+                w_value = from_ref(space, pvalue[0])
+                space.setitem(w_copy, w_key, w_value)
+        finally:
+            lltype.free(ppos, flavor='raw')
+            lltype.free(pkey, flavor='raw')
+            lltype.free(pvalue, flavor='raw')
+
+        api.Py_DecRef(py_dict) # release borrowed references
+
+        assert space.eq_w(space.len(w_copy), space.len(w_dict))
+        assert space.eq_w(w_copy, w_dict)