[pypy-commit] pypy ffi-backend: Test and fix: iteration

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: ffi-backend
Changeset: r55963:7e54c9509347
Date: 2012-07-07 13:27 +0200
http://bitbucket.org/pypy/pypy/changeset/7e54c9509347/

Log:	Test and fix: iteration

diff --git a/pypy/module/_cffi_backend/cdataobj.py b/pypy/module/_cffi_backend/cdataobj.py
--- a/pypy/module/_cffi_backend/cdataobj.py
+++ b/pypy/module/_cffi_backend/cdataobj.py
@@ -166,6 +166,9 @@
         keepalive_until_here(self)
         return w_result
 
+    def iter(self):
+        return self.ctype.iter(self)
+
     def write_raw_integer_data(self, source):
         misc.write_raw_integer_data(self._cdata, source, self.ctype.size)
         keepalive_until_here(self)
@@ -275,7 +278,8 @@
 
 
 W_CData.typedef = TypeDef(
-    '_cffi_backend.CData',
+    'CData',
+    __module__ = '_cffi_backend',
     __repr__ = interp2app(W_CData.repr),
     __nonzero__ = interp2app(W_CData.nonzero),
     __int__ = interp2app(W_CData.int),
@@ -293,12 +297,13 @@
     __getattr__ = interp2app(W_CData.getattr),
     __setattr__ = interp2app(W_CData.setattr),
     __call__ = interp2app(W_CData.call),
+    __iter__ = interp2app(W_CData.iter),
     )
 W_CData.typedef.acceptable_as_base_class = False
 
 W_CDataApplevelOwning.typedef = TypeDef(
-    '_cffi_backend.CDataOwn',
-    W_CData.typedef,    # base typedef
+    'CDataOwn', W_CData.typedef,    # base typedef
+    __module__ = '_cffi_backend',
     __weakref__ = make_weakref_descr(W_CDataApplevelOwning),
     )
 W_CDataApplevelOwning.typedef.acceptable_as_base_class = False
diff --git a/pypy/module/_cffi_backend/ctypearray.py b/pypy/module/_cffi_backend/ctypearray.py
--- a/pypy/module/_cffi_backend/ctypearray.py
+++ b/pypy/module/_cffi_backend/ctypearray.py
@@ -3,6 +3,9 @@
 """
 
 from pypy.interpreter.error import OperationError, operationerrfmt
+from pypy.interpreter.baseobjspace import Wrappable
+from pypy.interpreter.gateway import interp2app
+from pypy.interpreter.typedef import TypeDef
 from pypy.rpython.lltypesystem import rffi
 from pypy.rlib.objectmodel import keepalive_until_here
 from pypy.rlib.rarithmetic import ovfcheck
@@ -115,3 +118,35 @@
     def add(self, cdata, i):
         p = rffi.ptradd(cdata, i * self.ctitem.size)
         return cdataobj.W_CData(self.space, p, self.ctptr)
+
+    def iter(self, cdata):
+        return W_CDataIter(self.space, self.ctitem, cdata)
+
+
+class W_CDataIter(Wrappable):
+
+    def __init__(self, space, ctitem, cdata):
+        self.space = space
+        self.ctitem = ctitem
+        self.cdata = cdata
+        length = cdata.get_array_length()
+        self._next = cdata._cdata
+        self._stop = rffi.ptradd(cdata._cdata, length * ctitem.size)
+
+    def iter_w(self):
+        return self.space.wrap(self)
+
+    def next_w(self):
+        result = self._next
+        if result == self._stop:
+            raise OperationError(self.space.w_StopIteration, self.space.w_None)
+        self._next = rffi.ptradd(result, self.ctitem.size)
+        return self.ctitem.convert_to_object(result)
+
+W_CDataIter.typedef = TypeDef(
+    'CDataIter',
+    __module__ = '_cffi_backend',
+    __iter__ = interp2app(W_CDataIter.iter_w),
+    next = interp2app(W_CDataIter.next_w),
+    )
+W_CDataIter.typedef.acceptable_as_base_class = False
diff --git a/pypy/module/_cffi_backend/ctypeobj.py b/pypy/module/_cffi_backend/ctypeobj.py
--- a/pypy/module/_cffi_backend/ctypeobj.py
+++ b/pypy/module/_cffi_backend/ctypeobj.py
@@ -123,9 +123,16 @@
         raise operationerrfmt(space.w_TypeError,
                               "cdata '%s' is not callable", self.name)
 
+    def iter(self, cdata):
+        space = self.space
+        raise operationerrfmt(space.w_TypeError,
+                              "cdata '%s' does not support iteration",
+                              self.name)
+
 
 W_CType.typedef = TypeDef(
-    '_cffi_backend.CTypeDescr',
+    'CTypeDescr',
+    __module__ = '_cffi_backend',
     __repr__ = interp2app(W_CType.repr),
     __weakref__ = make_weakref_descr(W_CType),
     )
diff --git a/pypy/module/_cffi_backend/ctypestruct.py b/pypy/module/_cffi_backend/ctypestruct.py
--- a/pypy/module/_cffi_backend/ctypestruct.py
+++ b/pypy/module/_cffi_backend/ctypestruct.py
@@ -222,7 +222,8 @@
 
 
 W_CField.typedef = TypeDef(
-    '_cffi_backend.CField',
+    'CField',
+    __module__ = '_cffi_backend',
     type = interp_attrproperty('ctype', W_CField),
     offset = interp_attrproperty('offset', W_CField),
     bitshift = interp_attrproperty('bitshift', W_CField),
diff --git a/pypy/module/_cffi_backend/libraryobj.py b/pypy/module/_cffi_backend/libraryobj.py
--- a/pypy/module/_cffi_backend/libraryobj.py
+++ b/pypy/module/_cffi_backend/libraryobj.py
@@ -67,7 +67,8 @@
 
 
 W_Library.typedef = TypeDef(
-    '_cffi_backend.Library',
+    'Library',
+    __module__ = '_cffi_backend',
     __repr__ = interp2app(W_Library.repr),
     load_function = interp2app(W_Library.load_function),
     read_variable = interp2app(W_Library.read_variable),
diff --git a/pypy/module/_cffi_backend/test/_backend_test_c.py b/pypy/module/_cffi_backend/test/_backend_test_c.py
--- a/pypy/module/_cffi_backend/test/_backend_test_c.py
+++ b/pypy/module/_cffi_backend/test/_backend_test_c.py
@@ -1310,3 +1310,13 @@
     p = cast(BVoidP, 0)
     if 'PY_DOT_PY' in globals(): py.test.skip("NULL crashes early on py.py")
     py.test.raises(TypeError, "p[0]")
+
+def test_iter():
+    BInt = new_primitive_type("int")
+    BIntP = new_pointer_type(BInt)
+    BArray = new_array_type(BIntP, None)   # int[]
+    p = newp(BArray, 7)
+    assert list(p) == list(iter(p)) == [0] * 7
+    #
+    py.test.raises(TypeError, iter, cast(BInt, 5))
+    py.test.raises(TypeError, iter, cast(BIntP, 123456))