[pypy-commit] pypy fix-result-types: extract W_Ufunc1.find_specialization()

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: fix-result-types
Changeset: r77294:3ba5bdc3be93
Date: 2015-05-11 18:56 +0100
http://bitbucket.org/pypy/pypy/changeset/3ba5bdc3be93/

Log:	extract W_Ufunc1.find_specialization()

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
@@ -337,18 +337,7 @@
                 out = None
         w_obj = numpify(space, w_obj)
         dtype = w_obj.get_dtype(space)
-        if dtype.is_flexible():
-            raise OperationError(space.w_TypeError,
-                      space.wrap('Not implemented for this type'))
-        if (self.int_only and not (dtype.is_int() or dtype.is_object()) or
-                not self.allow_bool and dtype.is_bool() or
-                not self.allow_complex and dtype.is_complex()):
-            raise oefmt(space.w_TypeError,
-                "ufunc %s not supported for the input type", self.name)
-        calc_dtype = find_unaryop_result_dtype(space,
-                                  dtype,
-                                  promote_to_float=self.promote_to_float,
-                                  promote_bools=self.promote_bools)
+        calc_dtype, func = self.find_specialization(space, dtype)
         if out is not None:
             if not isinstance(out, W_NDimArray):
                 raise oefmt(space.w_TypeError, 'output must be an array')
@@ -371,7 +360,7 @@
         assert isinstance(w_obj, W_NDimArray)
         shape = shape_agreement(space, w_obj.get_shape(), out,
                                 broadcast_down=False)
-        return loop.call1(space, shape, self.func, calc_dtype, res_dtype,
+        return loop.call1(space, shape, func, calc_dtype, res_dtype,
                           w_obj, out)
 
     def call_scalar(self, space, w_arg, in_dtype, out_dtype, out):
@@ -387,6 +376,20 @@
             out.fill(space, w_val)
         return out
 
+    def find_specialization(self, space, dtype):
+        if dtype.is_flexible():
+            raise oefmt(space.w_TypeError, 'Not implemented for this type')
+        if (self.int_only and not (dtype.is_int() or dtype.is_object()) or
+                not self.allow_bool and dtype.is_bool() or
+                not self.allow_complex and dtype.is_complex()):
+            raise oefmt(space.w_TypeError,
+                "ufunc %s not supported for the input type", self.name)
+        calc_dtype = find_unaryop_result_dtype(space,
+                                  dtype,
+                                  promote_to_float=self.promote_to_float,
+                                  promote_bools=self.promote_bools)
+        return calc_dtype, self.func
+
 
 class W_Ufunc2(W_Ufunc):
     _immutable_fields_ = ["func", "comparison_func", "done_func"]