[pypy-commit] creflect default: ffi.addressof

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r173:832fea94ecfa
Date: 2014-12-05 18:40 +0100
http://bitbucket.org/cffi/creflect/changeset/832fea94ecfa/

Log:	ffi.addressof

diff --git a/zeffir/ffi_obj.c b/zeffir/ffi_obj.c
--- a/zeffir/ffi_obj.c
+++ b/zeffir/ffi_obj.c
@@ -317,6 +317,26 @@
     return NULL;
 }
 
+static CFieldObject *_ffi_field(CTypeDescrObject *ct, const char *fieldname)
+{
+    CFieldObject *cf;
+    if (ct->ct_stuff == NULL) {
+        PyErr_Format(PyExc_TypeError, "'%s' is incomplete", ct->ct_name);
+        return NULL;
+    }
+    cf = (CFieldObject *)PyDict_GetItemString(ct->ct_stuff, fieldname);
+    if (cf == NULL) {
+        PyErr_Format(PyExc_KeyError, "'%s' has got no field '%s'",
+                     ct->ct_name, fieldname);
+        return NULL;
+    }
+    if (cf->cf_bitshift >= 0) {
+        PyErr_SetString(PyExc_TypeError, "not supported for bitfields");
+        return NULL;
+    }
+    return cf;
+}
+
 static PyObject *ffi_offsetof(ZefFFIObject *self, PyObject *args)
 {
     PyObject *arg;
@@ -337,25 +357,47 @@
                      ct->ct_name);
         return NULL;
     }
-    if (ct->ct_stuff == NULL) {
-        PyErr_Format(PyExc_TypeError, "'%s' is incomplete", ct->ct_name);
+    cf = _ffi_field(ct, fieldname);
+    if (cf == NULL)
+        return NULL;
+    return PyInt_FromSsize_t(cf->cf_offset);
+}
+
+static PyObject *ffi_addressof(ZefFFIObject *self, PyObject *args)
+{
+    char *fieldname = NULL;
+    CDataObject *cd;
+    CTypeDescrObject *ct;
+    Py_ssize_t offset;
+
+    if (!PyArg_ParseTuple(args, "O!|z:addressof", &CData_Type, &cd, &fieldname))
+        return NULL;
+
+    ct = cd->c_type;
+    if (fieldname != NULL && ct->ct_flags & CT_POINTER)
+        ct = ct->ct_itemdescr;
+
+    if (!(ct->ct_flags & (CT_STRUCT|CT_UNION))) {
+        PyErr_Format(PyExc_TypeError,
+                     "expected a struct or union cdata, got '%s'",
+                     ct->ct_name);
         return NULL;
     }
 
-    cf = (CFieldObject *)PyDict_GetItemString(ct->ct_stuff, fieldname);
-    if (cf == NULL) {
-        PyErr_Format(PyExc_KeyError, "'%s' has got no field '%s'",
-                     ct->ct_name, fieldname);
-        return NULL;
+    if (fieldname == NULL) {
+        offset = 0;
     }
-    if (cf->cf_bitshift >= 0) {
-        PyErr_SetString(PyExc_TypeError, "not supported for bitfields");
-        return NULL;
+    else {
+        CFieldObject *cf = _ffi_field(ct, fieldname);
+        if (cf == NULL)
+            return NULL;
+        offset = cf->cf_offset;
     }
-    return PyInt_FromSsize_t(cf->cf_offset);
+    return new_simple_cdata(cd->c_data + offset, ct);
 }
 
 static PyMethodDef ffi_methods[] = {
+    {"addressof",     (PyCFunction)ffi_addressof,METH_VARARGS},
     {"close_library", ffi_close_library,         METH_VARARGS | METH_STATIC},
     {"load_library",  (PyCFunction)ffi_load_library,
                                                  METH_VARARGS | METH_KEYWORDS},
diff --git a/zeffir/test/test_struct.py b/zeffir/test/test_struct.py
--- a/zeffir/test/test_struct.py
+++ b/zeffir/test/test_struct.py
@@ -14,3 +14,15 @@
     ffi, lib = support.compile_and_open('struct')
     assert ffi.offsetof("mystruct_t", "a") == 0
     assert ffi.offsetof("mystruct_t", "b") == ffi.sizeof("long")
+
+def test_addressof():
+    ffi, lib = support.compile_and_open('struct')
+    p = ffi.new("mystruct_t *")
+    assert ffi.addressof(p[0]) == p
+    assert ffi.addressof(p[0], None) == p
+    assert ffi.addressof(p[0], "a") == p
+    assert ffi.addressof(p[0], "b") != p
+    assert ffi.addressof(p[0], "b") == ffi.addressof(p, "b")
+    assert ffi.addressof(p, "a") == p
+    py.test.raises(TypeError, ffi.addressof, p)
+    py.test.raises(TypeError, ffi.addressof, p, None)