Author: Brian Kearns <bdkea...@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 + + +@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') _______________________________________________ pypy-commit mailing list pypy-commit@python.org https://mail.python.org/mailman/listinfo/pypy-commit