[pypy-commit] pypy cpyext-injection: (arigato, plan_rich) hijack getitem of the minimal numpy module

[plan_rich]

Author: Richard Plangger <planrichi at gmail.com>
Branch: cpyext-injection
Changeset: r87893:fb5a66c1d78a
Date: 2016-10-20 15:07 +0200
http://bitbucket.org/pypy/pypy/changeset/fb5a66c1d78a/

Log:	(arigato, plan_rich) hijack getitem of the minimal numpy module

diff --git a/pypy/module/cpyext/injection/injection.py b/pypy/module/cpyext/injection/injection.py
--- a/pypy/module/cpyext/injection/injection.py
+++ b/pypy/module/cpyext/injection/injection.py
@@ -4,6 +4,10 @@
     if not we_are_translated() and name == 'test_module.test_mytype':
         from pypy.module.cpyext.injection._test_module import inject
         inject(space, name, dict_w, pto)
+    
+    if name == 'numpy.ndarray':
+        from pypy.module.cpyext.injection.numpy import inject_operator
+        inject_operator(space, name, dict_w, pto)
 
 def inject_global(space, w_func, modulename, funcname):
     if (not we_are_translated() and modulename == 'injection'
diff --git a/pypy/module/cpyext/injection/numpy.py b/pypy/module/cpyext/injection/numpy.py
--- a/pypy/module/cpyext/injection/numpy.py
+++ b/pypy/module/cpyext/injection/numpy.py
@@ -6,11 +6,12 @@
 from pypy.interpreter import typedef
 from pypy.objspace.std.intobject import W_IntObject
 from pypy.module.cpyext.api import bootstrap_function
+from pypy.interpreter.error import oefmt
 
 PyArrayObject = lltype.Ptr(lltype.Struct(
     'PyArrayObject',
     *(PyObjectFields + 
-        (("data", rffi.CHARP),
+        (("data", rffi.CCHARP),
          ("nd", rffi.INT),
          ("dimensions", rffi.SIGNEDP),
          ("strides", rffi.SIGNEDP),
@@ -39,3 +40,22 @@
     make_typedescr(W_ArrayObject.typedef,
                    basestruct=mytype_object.TO,
                    realize=mything_realize)
+
+ at unwrap_spec(index=int)
+def injected_getitem(space, w_self, index):
+    py_obj = rffi.cast(PyArrayObject, as_pyobj(space, w_self))
+    if index < 0 or index >= py_obj.dimensions[0]:
+        raise oefmt(space.w_IndexError, "index out of bounds")
+    data = rffi.cast(rffi.DOUBLEP, py_obj.data)
+    return space.newfloat(data[index])
+
+injected_methods = {
+    '__getitem__': interp2app(injected_getitem),
+}
+
+def inject_operator(space, name, dict_w, pto):
+    assert name == 'numpy.ndarray'
+    org = space.fromcache(Original)
+    org.w_original_getitem = dict_w['__getitem__']
+    for key, value in injected_methods.items():
+        dict_w[key] = space.wrap(value)
diff --git a/pypy/module/cpyext/injection/test/multiarray.c b/pypy/module/cpyext/injection/test/multiarray.c
--- a/pypy/module/cpyext/injection/test/multiarray.c
+++ b/pypy/module/cpyext/injection/test/multiarray.c
@@ -79,6 +79,10 @@
    }
    double * data = (double*)self->data;
    double value = data[i];
+   if (i == 10) {
+       // we try to modify this behaviour on the pypy level
+       value += 42;
+   }
 
    return PyFloat_FromDouble(value);
 }