[pypy-commit] pypy cpyext-injection: test one, add some tp_as_number methods

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: cpyext-injection
Changeset: r87954:ab7fe034304c
Date: 2016-10-26 23:36 +0300
http://bitbucket.org/pypy/pypy/changeset/ab7fe034304c/

Log:	test one, add some tp_as_number methods

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
@@ -1,4 +1,5 @@
 #include "Python.h"
+#include "pymath.h"
 
 typedef Py_intptr_t npy_intp;
 
@@ -275,6 +276,106 @@
     return NULL;
 }
 
+enum kOP {MULT, ADD, SUB, DIV};
+
+static PyObject*
+array_op(PyObject* obj1, PyObject* obj2, enum kOP op)
+{
+    int n1=-1, n2=-1, m, i1, i2, j;
+    double *v1, *v2, *r, tmp1, tmp2;
+    PyObject* ret = NULL, *tuple=NULL;
+    if (obj1->ob_type == &PyArray_Type)
+    {
+        n1 = ((PyArrayObject*)obj1)->dimensions[0];
+        v1 = (double*)((PyArrayObject*)obj1)->data;       
+    }
+    else
+    {
+        tmp1 = PyFloat_AsDouble(obj1);
+        if (PyErr_Occurred())
+            return NULL;
+        v1 = &tmp1;
+    }
+    if (obj2->ob_type == &PyArray_Type)
+    {
+        n2 = ((PyArrayObject*)obj2)->dimensions[0];
+        v2 = (double*)((PyArrayObject*)obj2)->data;       
+    }
+    else
+    {
+        tmp2 = PyFloat_AsDouble(obj2);
+        if (PyErr_Occurred())
+            return NULL;
+        v2 = &tmp2;
+    }
+    if (!(n1 == n2 || n1 == 1 || n2 == 1))
+    {
+        PyErr_SetString(PyExc_ValueError, "dimension mismatch");
+        return NULL;
+    }
+    m = n1 > n2? n1 : n2;
+    tuple = PyTuple_New(1);
+    PyTuple_SetItem(tuple, 0, PyInt_FromLong(m));
+    ret = array_new(&PyArray_Type, tuple, NULL);
+    Py_DECREF(tuple);
+    r = (double*)((PyArrayObject*)ret)->data;
+    for (i1=0, i2=0, j=0; j < m; j++, i1++, i2++)
+    {
+        if (i1 >= n1) i1 = 0;
+        if (i2 >= n2) i2 = 0;
+        switch (op)
+        {
+            case MULT:
+                r[j] = v1[i1] * v2[i2];
+                break;
+            case ADD:
+                r[j] = v1[i1] + v2[i2];
+                break;
+            case SUB:
+                r[j] = v1[i1] - v2[i2];
+                break;
+            case DIV:
+                if (v2[i2] == 0)
+                    r[j] = Py_NAN;
+                else
+                    r[j] = v1[i1] / v2[i2];
+                break;
+        }
+    } 
+    return ret;
+}
+
+static PyObject*
+array_multiply(PyObject* obj1, PyObject* obj2)
+{
+    return array_op(obj1, obj2, MULT);
+}
+
+static PyObject*
+array_add(PyObject* obj1, PyObject* obj2)
+{
+    return array_op(obj1, obj2, ADD);
+}
+
+static PyObject*
+array_sub(PyObject* obj1, PyObject* obj2)
+{
+    return array_op(obj1, obj2, SUB);
+}
+
+static PyObject*
+array_divide(PyObject* obj1, PyObject* obj2)
+{
+    return array_op(obj1, obj2, DIV);
+}
+
+static PyNumberMethods array_as_number = {
+    (binaryfunc)array_add, /* nb_add*/
+    (binaryfunc)array_sub, /* nb_subtract */
+    (binaryfunc)array_multiply, /* nb_multiply */
+    (binaryfunc)array_divide, /* nb_divide */
+};
+
 /* List of functions exported by this module */
 
 static PyMethodDef multiarray_functions[] = {
@@ -330,6 +431,8 @@
     if (module == NULL)
         INITERROR;
 
+    PyArray_Type.tp_as_number = &array_as_number;
+    
     if (PyType_Ready(&PyArray_Type) < 0)
         INITERROR;
     PyModule_AddObject(module, "ndarray", (PyObject *)&PyArray_Type);
diff --git a/pypy/module/cpyext/injection/test/test_numpy.py b/pypy/module/cpyext/injection/test/test_numpy.py
--- a/pypy/module/cpyext/injection/test/test_numpy.py
+++ b/pypy/module/cpyext/injection/test/test_numpy.py
@@ -36,3 +36,12 @@
         x = X(1)
         assert isinstance(x, X)
         assert isinstance(x, np.ndarray)
+
+    def test_plain_op(self):
+        np = self.import_module(name='numpy.core.multiarray_PLAIN',
+                                filename='../injection/test/multiarray')
+        a = np.ndarray(100);
+        for i in range(100):
+            a[i] = i
+        b = a * a
+        assert b[10] == 100.0 + 42.0