[pypy-commit] cffi default: Test iteration over array cdatas, and found out that we need more code.

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r402:5a64a124c201
Date: 2012-06-17 14:46 +0200
http://bitbucket.org/cffi/cffi/changeset/5a64a124c201/

Log:	Test iteration over array cdatas, and found out that we need more
	code.

diff --git a/c/_ffi_backend.c b/c/_ffi_backend.c
--- a/c/_ffi_backend.c
+++ b/c/_ffi_backend.c
@@ -1497,6 +1497,8 @@
     goto done;
 }
 
+static PyObject *cdata_iter(CDataObject *);
+
 static PyNumberMethods CData_as_number = {
     (binaryfunc)cdata_add,      /*nb_add*/
     (binaryfunc)cdata_sub,      /*nb_subtract*/
@@ -1554,6 +1556,8 @@
     (traverseproc)cdata_traverse,               /* tp_traverse */
     0,                                          /* tp_clear */
     cdata_richcompare,                          /* tp_richcompare */
+    0,                                          /* tp_weaklistoffset */
+    (getiterfunc)cdata_iter,                    /* tp_iter */
 };
 
 static PyTypeObject CDataOwning_Type = {
@@ -1624,6 +1628,89 @@
     &CData_Type,                                /* tp_base */
 };
 
+/************************************************************/
+
+typedef struct {
+    PyObject_HEAD
+    char *di_next, *di_stop;
+    CDataObject *di_object;
+    CTypeDescrObject *di_itemtype;
+} CDataIterObject;
+
+static PyObject *
+cdataiter_next(CDataIterObject *it)
+{
+    char *result = it->di_next;
+    if (result != it->di_stop) {
+        it->di_next = result + it->di_itemtype->ct_size;
+        return convert_to_object(result, it->di_itemtype);
+    }
+    return NULL;
+}
+
+static void
+cdataiter_dealloc(CDataIterObject *it)
+{
+    Py_DECREF(it->di_object);
+    PyObject_Del(it);
+}
+
+static PyTypeObject CDataIter_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    "_ffi_backend.CDataIter",               /* tp_name */
+    sizeof(CDataIterObject),                /* tp_basicsize */
+    0,                                      /* tp_itemsize */
+    /* methods */
+    (destructor)cdataiter_dealloc,          /* tp_dealloc */
+    0,                                      /* tp_print */
+    0,                                      /* tp_getattr */
+    0,                                      /* tp_setattr */
+    0,                                      /* tp_compare */
+    0,                                      /* tp_repr */
+    0,                                      /* tp_as_number */
+    0,                                      /* tp_as_sequence */
+    0,                                      /* tp_as_mapping */
+    0,                                      /* tp_hash */
+    0,                                      /* tp_call */
+    0,                                      /* tp_str */
+    PyObject_GenericGetAttr,                /* tp_getattro */
+    0,                                      /* tp_setattro */
+    0,                                      /* tp_as_buffer */
+    Py_TPFLAGS_DEFAULT,                     /* tp_flags */
+    0,                                      /* tp_doc */
+    0,                                      /* tp_traverse */
+    0,                                      /* tp_clear */
+    0,                                      /* tp_richcompare */
+    0,                                      /* tp_weaklistoffset */
+    PyObject_SelfIter,                      /* tp_iter */
+    (iternextfunc)cdataiter_next,           /* tp_iternext */
+};
+
+static PyObject *
+cdata_iter(CDataObject *cd)
+{
+    CDataIterObject *it;
+
+    if (!(cd->c_type->ct_flags & CT_ARRAY)) {
+        PyErr_Format(PyExc_TypeError, "cdata '%s' does not support iteration",
+                     cd->c_type->ct_name);
+        return NULL;
+    }
+
+    it = PyObject_New(CDataIterObject, &CDataIter_Type);
+    if (it == NULL)
+        return NULL;
+
+    Py_INCREF(cd);
+    it->di_object = cd;
+    it->di_itemtype = cd->c_type->ct_itemdescr;
+    it->di_next = cd->c_data;
+    it->di_stop = cd->c_data + get_array_length(cd) * it->di_itemtype->ct_size;
+    return (PyObject *)it;
+}
+
+/************************************************************/
+
 static PyObject *b_newp(PyObject *self, PyObject *args)
 {
     CTypeDescrObject *ct, *ctitem;
@@ -3443,6 +3530,8 @@
         return;
     if (PyType_Ready(&CDataWithDestructor_Type) < 0)
         return;
+    if (PyType_Ready(&CDataIter_Type) < 0)
+        return;
 
     v = PyCObject_FromVoidPtr((void *)cffi_exports, NULL);
     if (v == NULL || PyModule_AddObject(m, "_C_API", v) < 0)
diff --git a/cffi/backend_ctypes.py b/cffi/backend_ctypes.py
--- a/cffi/backend_ctypes.py
+++ b/cffi/backend_ctypes.py
@@ -72,6 +72,10 @@
     def _alignment(cls):
         return ctypes.alignment(cls._ctype)
 
+    def __iter__(self):
+        raise TypeError("cdata %r does not support iteration" % (
+            self._get_c_name()),)
+
 
 class CTypesGenericPrimitive(CTypesData):
     __slots__ = []
@@ -80,6 +84,10 @@
 class CTypesGenericArray(CTypesData):
     __slots__ = []
 
+    def __iter__(self):
+        for i in xrange(len(self)):
+            yield self[i]
+
 
 class CTypesGenericPtr(CTypesData):
     __slots__ = ['_address', '_as_ctype_ptr']
diff --git a/testing/backend_tests.py b/testing/backend_tests.py
--- a/testing/backend_tests.py
+++ b/testing/backend_tests.py
@@ -694,6 +694,17 @@
         p = ffi.new("int(*)[5]")
         assert repr(p) == "<cdata 'int(* *)[5]' owning %d bytes>" % SIZE_OF_PTR
 
+    def test_iterate_array(self):
+        ffi = FFI(backend=self.Backend())
+        a = ffi.new("char[]", "hello")
+        assert list(a) == ["h", "e", "l", "l", "o", chr(0)]
+        assert list(iter(a)) == ["h", "e", "l", "l", "o", chr(0)]
+        #
+        py.test.raises(TypeError, iter, ffi.cast("char *", a))
+        py.test.raises(TypeError, list, ffi.cast("char *", a))
+        py.test.raises(TypeError, iter, ffi.new("int"))
+        py.test.raises(TypeError, list, ffi.new("int"))
+
     def test_offsetof(self):
         ffi = FFI(backend=self.Backend())
         ffi.cdef("struct foo { int a, b, c; };")