[pypy-commit] pypy cpyext-injection: Injecting a module's global function

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: cpyext-injection
Changeset: r87874:ab5a54a34d42
Date: 2016-10-19 15:45 +0200
http://bitbucket.org/pypy/pypy/changeset/ab5a54a34d42/

Log:	Injecting a module's global function

diff --git a/pypy/module/cpyext/injection/_test_module.py b/pypy/module/cpyext/injection/_test_module.py
--- a/pypy/module/cpyext/injection/_test_module.py
+++ b/pypy/module/cpyext/injection/_test_module.py
@@ -26,9 +26,15 @@
         return space.call_function(org.w_original_getitem, w_self,
                                    space.wrap(index))
 
+ at unwrap_spec(arg=int)
+def injected_make(space, arg):
+    return space.w_Ellipsis
+
+
 injected_methods = {
     '__getitem__': interp2app(injected_getitem),
 }
+interp_injected_make = interp2app(injected_make)
 
 def inject(space, name, dict_w, pto):
     assert name == 'test_module.test_mytype'
@@ -36,3 +42,9 @@
     org.w_original_getitem = dict_w['__getitem__']
     for key, value in injected_methods.items():
         dict_w[key] = space.wrap(value)
+
+def inject_global(space, w_func, name):
+    assert name == 'make'
+    org = space.fromcache(Original)
+    org.w_original_make = w_func
+    return space.wrap(interp_injected_make)
diff --git a/pypy/module/cpyext/modsupport.py b/pypy/module/cpyext/modsupport.py
--- a/pypy/module/cpyext/modsupport.py
+++ b/pypy/module/cpyext/modsupport.py
@@ -1,4 +1,5 @@
 from rpython.rtyper.lltypesystem import rffi, lltype
+from rpython.rlib.objectmodel import we_are_translated
 from pypy.module.cpyext.api import cpython_api, cpython_struct, \
         METH_STATIC, METH_CLASS, METH_COEXIST, CANNOT_FAIL, CONST_STRING
 from pypy.module.cpyext.pyobject import PyObject
@@ -91,6 +92,7 @@
                             "module functions cannot set METH_CLASS or "
                             "METH_STATIC")
                 w_obj = space.wrap(W_PyCFunctionObject(space, method, w_self, w_name))
+                w_obj = inject_global(space, w_obj, name, methodname)
             else:
                 if methodname in dict_w and not (flags & METH_COEXIST):
                     continue
@@ -132,3 +134,9 @@
     raise NotImplementedError
 
 
+def inject_global(space, w_func, modulename, funcname):
+    if (not we_are_translated() and modulename == 'injection'
+          and funcname == 'make'):
+        from pypy.module.cpyext.injection._test_module import inject_global
+        w_func = inject_global(space, w_func, funcname)
+    return w_func
diff --git a/pypy/module/cpyext/test/injection.c b/pypy/module/cpyext/test/injection.c
--- a/pypy/module/cpyext/test/injection.c
+++ b/pypy/module/cpyext/test/injection.c
@@ -70,9 +70,23 @@
     PyObject_Del,                               /* tp_free */
 };
 
+
+static PyObject *glob_make(PyObject *self, PyObject *args)
+{
+    int i;
+    if (!PyArg_ParseTuple(args, "i", &i))
+        return NULL;
+
+    PyTypeObject *type = &mytype_type;
+    mytype_object *o = (mytype_object *)type->tp_alloc(type, 0);
+    o->foo = i;
+    return (PyObject *)o;
+}
+
 /* List of functions exported by this module */
 
 static PyMethodDef foo_functions[] = {
+    {"make",      (PyCFunction)glob_make, METH_VARARGS, NULL},
     {NULL,        NULL}    /* Sentinel */
 };
 
diff --git a/pypy/module/cpyext/test/test_injection.py b/pypy/module/cpyext/test/test_injection.py
--- a/pypy/module/cpyext/test/test_injection.py
+++ b/pypy/module/cpyext/test/test_injection.py
@@ -8,3 +8,8 @@
         mything = module.test_mytype()
         assert mything[100] == 4200
         assert mything[-100] == -100+42
+
+    def test_glob_make(self):
+        module = self.import_module(name='injection')
+        mything = module.make(5)
+        assert mything is Ellipsis