[pypy-commit] pypy ufuncapi: do enough parsing, start iterator creation

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: ufuncapi
Changeset: r73994:0acd1768b6b6
Date: 2014-10-16 08:27 -0500
http://bitbucket.org/pypy/pypy/changeset/0acd1768b6b6/

Log:	do enough parsing, start iterator creation

diff --git a/pypy/module/micronumpy/nditer.py b/pypy/module/micronumpy/nditer.py
--- a/pypy/module/micronumpy/nditer.py
+++ b/pypy/module/micronumpy/nditer.py
@@ -347,6 +347,8 @@
         return space.wrap(self)
 
     def getitem(self, it, st, op_flags):
+        # TODO: push op_flags and the array onto the iterator
+        # so we can avoid creating a new W_NDimArray each iteration
         if op_flags.rw == 'r':
             impl = concrete.ConcreteNonWritableArrayWithBase
         else:
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -160,6 +160,21 @@
         af2 = ufunc(af)
         assert all(af2 == af * 2)
 
+    def test_frompyfunc_2d_sig(self):
+        def times_2(in_array, out_array):
+            in_flat = in_array.flat
+            out_flat = out_array.flat
+            for i in range(in_array.size):
+                out_flat[i] = in_flat[i] * 2
+        from numpy import frompyfunc, dtype, arange
+        ufunc = frompyfunc([times_2], 1, 1,
+                            signature='(m,m)->(m,m)',
+                            dtypes=[dtype(int), dtype(int)],
+                          )
+        ai = arange(18, dtype=int).reshape(2,3,3)
+        ai2 = ufunc(ai)
+        assert all(ai1 == ai * 2)
+
     def test_ufunc_kwargs(self):
         from numpy import ufunc, frompyfunc, arange, dtype
         def adder(a, b):
diff --git a/pypy/module/micronumpy/ufuncs.py b/pypy/module/micronumpy/ufuncs.py
--- a/pypy/module/micronumpy/ufuncs.py
+++ b/pypy/module/micronumpy/ufuncs.py
@@ -580,8 +580,6 @@
         if not self.core_enabled:
             # func is going to do all the work, it must accept W_NDimArray args
             arglist = space.newlist(list(inargs + outargs))
-            if not isinstance(func, W_GenericUFuncCaller):
-                raise oefmt(space.w_RuntimeError, "cannot do core_enabled frompyfunc")
             space.call_args(func, Arguments.frompacked(space, arglist))
             if len(outargs) < 2:
                 return outargs[0]
@@ -607,7 +605,7 @@
         # into the array 'inner_dimensions[1:]'. Initialize them to
         # 1, for example in the case where the operand broadcasts
         # to a core dimension, it won't be visited.
-        inner_dimensions = [1] * (self.core_num_dim_ix + 1)
+        inner_dimensions = [1] * (self.core_num_dim_ix + 2)
         idim = 0
         for i in range(self.nin):
             curarg = inargs[i]
@@ -656,7 +654,7 @@
                         "(size %d is different from %d)", self.name, i, idim,
                     self.signature, op_dim_size, inner_dimensions[1 + core_dim_index])
                 idim += 1
-        iter_shape = [-1] * (broadcast_ndim + len(outargs))
+        iter_shape = [-1] * (broadcast_ndim + (len(outargs) * num_dims))
         j = broadcast_ndim
         core_dim_ixs_size = 0
         firstdim = broadcast_ndim
@@ -681,13 +679,13 @@
             for idim in range(broadcast_ndim):
                 if idim >= broadcast_ndim -n:
                     op_axes_arrays[idim][iout] = idim - (broadcast_ndim -n)
-                dim_offset = self.core_offsets[iout]
-                num_dims = self.core_num_dims[iout]
-                for idim in range(num_dims):
-                    cdi = self.core_dim_ixs[dim_offset + idim]
-                    iter_shape[j] = inner_dimensions[1 + cdi]
-                    op_axes_arrays[j][iout] = n + idim
-                    j += 1
+            dim_offset = self.core_offsets[iout]
+            num_dims = self.core_num_dims[iout]
+            for idim in range(num_dims):
+                cdi = self.core_dim_ixs[dim_offset + idim]
+                iter_shape[j] = inner_dimensions[1 + cdi]
+                op_axes_arrays[j][iout] = n + idim
+                j += 1
             core_dim_ixs_size += self.core_num_dims[iout];
         # TODO once we support obejct dtypes,
         # FAIL with NotImplemented if the other object has
@@ -696,6 +694,8 @@
 
         # TODO parse and handle subok
 
+        # mimic NpyIter_AdvancedNew with a nditer
+
         if isinstance(func, W_GenericUFuncCaller):
             pass
             # xxx rdo what needs to be done to inner-loop indexing