more numpy arrays conversions

fijal · Dec 2, 2014 · 693f562 · 693f562
1 parent f1ad0a4
commit 693f562
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 3 deletions.
diff --git a/jitpy/__init__.py b/jitpy/__init__.py
@@ -55,6 +55,9 @@ def setup(pypy_home):
     } PyObject;
 
     typedef struct {
+       char kind;
+       char type;
+       int type_num;
        ...;
     } PyArray_Descr;
     
@@ -75,11 +78,13 @@ def setup(pypy_home):
     #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
     #include <numpy/arrayobject.h>
     """)
-    ptr.setup_numpy_data(ffi2.new("int[4]", [
+    ptr.setup_numpy_data(ffi2.new("int[6]", [
         ffi2.offsetof("PyArrayObject_fields", "data"),
         ffi2.offsetof("PyArrayObject_fields", "strides"),
         ffi2.offsetof("PyArrayObject_fields", "dimensions"),
         ffi2.offsetof("PyArrayObject_fields", "nd"),
+        ffi2.offsetof("PyArrayObject_fields", "descr"),
+        ffi2.offsetof("PyArray_Descr", "type_num"),
     ]))
 
     globals()['ffi'] = ffi

diff --git a/jitpy/pypy_side.py b/jitpy/pypy_side.py
@@ -16,6 +16,29 @@
 else:
     WORD = 8
 
+# XXX copy
+numpy_types = {
+    0: 'b',
+    1: 'i1',
+    2: 'u1',
+    3: 'i2',
+    4: 'u2',
+    5: 'int',
+    6: 'uint',
+    7: 'int',
+    8: 'uint',
+    9: 'longlong',
+    10: 'ulonglong',
+    11: 'single',
+    12: 'double',
+    13: 'longdouble',
+    14: 'csingle',
+    15: 'cdouble',
+    16: 'clongdouble',
+    #17: 'object',
+    # rest not supported
+}
+
 def convert_from_numpy_array(ll_a):
     data = (ll_a + NumpyData.data_offset)[0]
     strides = ffi.cast("intptr_t*", (ll_a + NumpyData.strides_offset)[0])
@@ -25,7 +48,10 @@ def convert_from_numpy_array(ll_a):
     strides = [strides[i] for i in range(nd)]
     data = ffi.buffer(ffi.cast("char*", data), sys.maxint)
     # XXX strides
-    return numpy.ndarray(dims, buffer=data, dtype=int)
+    descr = ffi.cast('char*', (ll_a + NumpyData.descr_offset)[0])
+    type_num = ffi.cast('int*', descr + NumpyData.type_num_offset)[0]
+    dtype = numpy_types[type_num]
+    return numpy.ndarray(dims, buffer=data, dtype=dtype)
 
 def wrap_exception_and_arrays(func, arrays):
     def wrapper(*args):
@@ -77,6 +103,8 @@ def setup_numpy_data(offsets):
     NumpyData.strides_offset = offsets[1] / WORD
     NumpyData.dimensions_offset = offsets[2] / WORD
     NumpyData.nd_offset = offsets[3] / WORD
+    NumpyData.descr_offset = offsets[4] / WORD
+    NumpyData.type_num_offset = offsets[5]
 
 pypy_def = ffi.cast("struct pypy_defs*", c_argument)
 pypy_def.basic_register = basic_register

diff --git a/jitpy/wrapper.py b/jitpy/wrapper.py
@@ -29,7 +29,6 @@ def jittify(argtypes, restype):
 
     def convert(ll_tp, arg):
         if ll_tp == 'array':
-            assert arg.dtype == int # for now
             return ffi.cast("void *", id(arg))
         return arg
 

diff --git a/testing/test_jitpy.py b/testing/test_jitpy.py
@@ -56,3 +56,27 @@ def func(a, s):
         a = numpy.array([1, 2, 3])
         func(a, 3)
         assert a[0] == 4
+
+    def test_numpy_array_float(self):
+        import numpy
+
+        @jittify(['array', int], None)
+        def func(a, s):
+            for i in range(a.shape[0]):
+                a[i] += s
+
+        a = numpy.array([1.2, 2, 3], dtype=float)
+        func(a, 3)
+        assert a[0] == 1.2 + 3
+
+    def test_numpy_array_singlefloat(self):
+        import numpy
+
+        @jittify(['array', int], None)
+        def func(a, s):
+            for i in range(a.shape[0]):
+                a[i] += s
+
+        a = numpy.array([1.2, 2, 3], dtype='f32')
+        func(a, 3)
+        assert (a[0] - (1.2 + 3)) < 0.0000001 # inexact