[pypy-commit] pypy numpypy-array_prepare_-array_wrap: allow Box results to percolate through, start to fix compile.py

[mattip]

Author: Matti Picus <matti.picus at gmail.com>
Branch: numpypy-array_prepare_-array_wrap
Changeset: r67040:e69c4253c9d7
Date: 2013-09-20 16:28 +0300
http://bitbucket.org/pypy/pypy/changeset/e69c4253c9d7/

Log:	allow Box results to percolate through, start to fix compile.py

diff --git a/pypy/module/micronumpy/compile.py b/pypy/module/micronumpy/compile.py
--- a/pypy/module/micronumpy/compile.py
+++ b/pypy/module/micronumpy/compile.py
@@ -222,6 +222,10 @@
         return None
         #return getattr(w_obj, 'descr_' + s)(self, *args)
 
+    def get_and_call_function(self, w_descr, w_obj, *args_w):
+        w_impl = self.getattr(w_descr, w_obj)
+        return self.call_method(w_obj, w_impl, args_w)
+
     @specialize.arg(1)
     def interp_w(self, tp, what):
         assert isinstance(what, tp)
diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -439,7 +439,7 @@
 
     def descr___array_prepare__(self, space, w_array, w_context):
         # stub implementation of __array_prepare__()
-        if isinstance(w_array, W_NDimArray):
+        if isinstance(w_array, (W_NDimArray, interp_boxes.Box)):
             return w_array
         else:
             raise OperationError(space.w_TypeError,
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -234,6 +234,37 @@
             return out
         return res
 
+    def call_prepare(self, space, w_out, w_obj, w_result):
+        if isinstance(w_out, W_NDimArray):
+            w_array = space.lookup(w_out, "__array_prepare__")
+            w_caller = w_out
+        else:
+            w_array = space.lookup(w_obj, "__array_prepare__")
+            w_caller = w_obj
+        if w_array:
+            w_retVal = space.get_and_call_function(w_array, w_caller, w_result, None)
+            if not isinstance(w_retVal, (W_NDimArray, interp_boxes.Box)):
+                raise OperationError(space.w_ValueError,
+                    space.wrap( "__array_prepare__ must return an "
+                                "ndarray or subclass thereof"))
+            if isinstance(w_result, interp_boxes.Box) or \
+                w_result.is_scalar():
+                if not isinstance(w_retVal, interp_boxes.Box) and not w_retVal.is_scalar():
+                    raise OperationError(space.w_TypeError,
+                        space.wrap( "__array_prepare__ must return an "
+                                    "ndarray or subclass thereof which is "
+                                    "otherwise identical to its input"))
+            elif w_result.get_shape() != w_retVal.get_shape() or \
+               w_result.implementation.get_strides() != \
+                    w_retVal.implementation.get_strides():
+                raise OperationError(space.w_TypeError,
+                    space.wrap( "__array_prepare__ must return an "
+                                "ndarray or subclass thereof which is "
+                                "otherwise identical to its input"))
+            return w_retVal
+        return w_result
+
+
 class W_Ufunc1(W_Ufunc):
     argcount = 1
 
@@ -319,22 +350,6 @@
         else:
             self.done_func = None
 
-    def call_prepare(self, space, w_out, w_obj, w_result):
-        if isinstance(w_out, W_NDimArray):
-            w_array = space.lookup(w_out, "__array_prepare__")
-            w_caller = w_out
-        else:
-            w_array = space.lookup(w_obj, "__array_prepare__")
-            w_caller = w_obj
-        if w_array:
-            w_result = space.get_and_call_function(w_array, w_caller, w_result, None)
-            if not isinstance(w_result, W_NDimArray):
-                raise OperationError(space.w_ValueError,
-                        space.wrap("object __array_prepare__ method not"
-                                   " producing an array"))
-        return w_result
-
-
     @jit.unroll_safe
     def call(self, space, args_w):
         if len(args_w) > 2:
@@ -387,6 +402,7 @@
                     out.set_scalar_value(arr)
                 else:
                     out.fill(arr)
+                arr = out
             return self.call_prepare(space, out, w_lhs, arr)
         new_shape = shape_agreement(space, w_lhs.get_shape(), w_rhs)
         new_shape = shape_agreement(space, new_shape, out, broadcast_down=False)