[pypy-commit] pypy default: implement numpy.result_type function

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r73865:84e2e30e7551
Date: 2014-10-09 12:40 -0400
http://bitbucket.org/pypy/pypy/changeset/84e2e30e7551/

Log:	implement numpy.result_type function

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -20,6 +20,7 @@
         'concatenate': 'arrayops.concatenate',
         'count_nonzero': 'arrayops.count_nonzero',
         'dot': 'arrayops.dot',
+        'result_type': 'arrayops.result_type',
         'where': 'arrayops.where',
 
         'set_string_function': 'appbridge.set_string_function',
diff --git a/pypy/module/micronumpy/arrayops.py b/pypy/module/micronumpy/arrayops.py
--- a/pypy/module/micronumpy/arrayops.py
+++ b/pypy/module/micronumpy/arrayops.py
@@ -1,3 +1,4 @@
+from rpython.rlib import jit
 from pypy.interpreter.error import OperationError, oefmt
 from pypy.interpreter.gateway import unwrap_spec
 from pypy.module.micronumpy import loop, descriptor, ufuncs, support, \
@@ -6,6 +7,7 @@
 from pypy.module.micronumpy.converters import clipmode_converter
 from pypy.module.micronumpy.strides import Chunk, Chunks, shape_agreement, \
     shape_agreement_multiple
+from .boxes import W_GenericBox
 
 
 def where(space, w_arr, w_x=None, w_y=None):
@@ -283,3 +285,26 @@
     else:
         loop.diagonal_array(space, arr, out, offset, axis1, axis2, shape)
     return out
+
+
+ at jit.unroll_safe
+def result_type(space, __args__):
+    args_w, kw_w = __args__.unpack()
+    if kw_w:
+        raise oefmt(space.w_TypeError, "result_type() takes no keyword arguments")
+    result = None
+    for w_arg in args_w:
+        if isinstance(w_arg, W_NDimArray):
+            dtype = w_arg.get_dtype()
+        elif isinstance(w_arg, W_GenericBox) or (
+                space.isinstance_w(w_arg, space.w_int) or
+                space.isinstance_w(w_arg, space.w_float) or
+                space.isinstance_w(w_arg, space.w_complex) or
+                space.isinstance_w(w_arg, space.w_long) or
+                space.isinstance_w(w_arg, space.w_bool)):
+            dtype = ufuncs.find_dtype_for_scalar(space, w_arg)
+        else:
+            dtype = space.interp_w(descriptor.W_Dtype,
+                space.call_function(space.gettypefor(descriptor.W_Dtype), w_arg))
+        result = ufuncs.find_binop_result_dtype(space, result, dtype)
+    return result
diff --git a/pypy/module/micronumpy/test/test_arrayops.py b/pypy/module/micronumpy/test/test_arrayops.py
--- a/pypy/module/micronumpy/test/test_arrayops.py
+++ b/pypy/module/micronumpy/test/test_arrayops.py
@@ -199,3 +199,17 @@
         a.put(23, -1, mode=1)  # wrap
         assert (a == array([0, 1, -10, -1, -15])).all()
         raises(TypeError, "arange(5).put(22, -5, mode='zzzz')")  # unrecognized mode
+
+    def test_result_type(self):
+        import numpy as np
+        exc = raises(TypeError, np.result_type, a=2)
+        assert str(exc.value) == "result_type() takes no keyword arguments"
+        assert np.result_type(True) is np.dtype('bool')
+        assert np.result_type(1) is np.dtype('int64')
+        assert np.result_type(1.) is np.dtype('float64')
+        assert np.result_type(1+2j) is np.dtype('complex128')
+        assert np.result_type(1, 1.) is np.dtype('float64')
+        assert np.result_type(np.array([1, 2])) is np.dtype('int64')
+        assert np.result_type(np.array([1, 2]), 1, 1+2j) is np.dtype('complex128')
+        assert np.result_type(np.array([1, 2]), 1, 'float64') is np.dtype('float64')
+        assert np.result_type(np.array([1, 2]), 1, None) is np.dtype('float64')