"""Functions and helpers for converting between dtypes""" from rpython.rlib import jit, objectmodel from rpython.rlib.signature import signature, types as ann from pypy.interpreter.gateway import unwrap_spec from pypy.interpreter.error import OperationError, oefmt from pypy.module.micronumpy.base import W_NDimArray, convert_to_array from pypy.module.micronumpy import constants as NPY from .types import ( BaseType, Bool, ULong, Long, Float64, Complex64, StringType, UnicodeType, VoidType, ObjectType, int_types, float_types, complex_types, number_types, all_types) from .descriptor import ( W_Dtype, get_dtype_cache, as_dtype, is_scalar_w, variable_dtype, new_string_dtype, new_unicode_dtype, num2dtype) @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") if not args_w: raise oefmt(space.w_ValueError, "at least one array or dtype is required") arrays_w = [] dtypes_w = [] for w_arg in args_w: if isinstance(w_arg, W_NDimArray): arrays_w.append(w_arg) elif is_scalar_w(space, w_arg): w_scalar = as_scalar(space, w_arg) w_arr = W_NDimArray.from_scalar(space, w_scalar) arrays_w.append(w_arr) else: dtype = as_dtype(space, w_arg) dtypes_w.append(dtype) return find_result_type(space, arrays_w, dtypes_w) @jit.look_inside_iff(lambda space, arrays_w, dtypes_w: jit.loop_unrolling_heuristic(arrays_w, len(arrays_w)) and jit.loop_unrolling_heuristic(dtypes_w, len(dtypes_w))) def find_result_type(space, arrays_w, dtypes_w): # equivalent to PyArray_ResultType if len(arrays_w) == 1 and not dtypes_w: return arrays_w[0].get_dtype() elif not arrays_w and len(dtypes_w) == 1: return dtypes_w[0] result = None if not _use_min_scalar(arrays_w, dtypes_w): for w_array in arrays_w: if result is None: result = w_array.get_dtype() else: result = promote_types(space, result, w_array.get_dtype()) for dtype in dtypes_w: if result is None: result = dtype else: result = promote_types(space, result, dtype) else: small_unsigned = False for w_array in arrays_w: dtype = w_array.get_dtype() small_unsigned_scalar = False if w_array.is_scalar() and dtype.is_number(): num, alt_num = w_array.get_scalar_value().min_dtype() small_unsigned_scalar = (num != alt_num) dtype = num2dtype(space, num) if result is None: result = dtype small_unsigned = small_unsigned_scalar else: result, small_unsigned = _promote_types_su( space, result, dtype, small_unsigned, small_unsigned_scalar) for dtype in dtypes_w: if result is None: result = dtype small_unsigned = False else: result, small_unsigned = _promote_types_su( space, result, dtype, small_unsigned, False) return result simple_kind_ordering = objectmodel.dict_to_switch({ Bool.kind: 0, ULong.kind: 1, Long.kind: 1, Float64.kind: 2, Complex64.kind: 2, NPY.STRINGLTR: 3, NPY.STRINGLTR2: 3, UnicodeType.kind: 3, VoidType.kind: 3, ObjectType.kind: 3}) # this is safe to unroll since it'll only be seen if we look inside # the find_result_type @jit.unroll_safe def _use_min_scalar(arrays_w, dtypes_w): """Helper for find_result_type()""" if not arrays_w: return False all_scalars = True max_scalar_kind = 0 max_array_kind = 0 for w_array in arrays_w: if w_array.is_scalar(): kind = simple_kind_ordering(w_array.get_dtype().kind) if kind > max_scalar_kind: max_scalar_kind = kind else: all_scalars = False kind = simple_kind_ordering(w_array.get_dtype().kind) if kind > max_array_kind: max_array_kind = kind for dtype in dtypes_w: all_scalars = False kind = simple_kind_ordering(dtype.kind) if kind > max_array_kind: max_array_kind = kind return not all_scalars and max_array_kind >= max_scalar_kind @unwrap_spec(casting='text') def can_cast(space, w_from, w_totype, casting='safe'): try: target = as_dtype(space, w_totype, allow_None=False) except TypeError: raise oefmt(space.w_TypeError, "did not understand one of the types; 'None' not accepted") if isinstance(w_from, W_NDimArray): return space.newbool(can_cast_array(space, w_from, target, casting)) elif is_scalar_w(space, w_from): w_scalar = as_scalar(space, w_from) w_arr = W_NDimArray.from_scalar(space, w_scalar) return space.newbool(can_cast_array(space, w_arr, target, casting)) try: origin = as_dtype(space, w_from, allow_None=False) except TypeError: raise oefmt(space.w_TypeError, "did not understand one of the types; 'None' not accepted") return space.newbool(can_cast_type(space, origin, target, casting)) kind_ordering = { Bool.kind: 0, ULong.kind: 1, Long.kind: 2, Float64.kind: 4, Complex64.kind: 5, NPY.STRINGLTR: 6, NPY.STRINGLTR2: 6, UnicodeType.kind: 7, VoidType.kind: 8, ObjectType.kind: 9} def can_cast_type(space, origin, target, casting): # equivalent to PyArray_CanCastTypeTo if origin == target: return True if casting == 'unsafe': return True elif casting == 'no': return origin.eq(space, target) if origin.num == target.num: if origin.is_record(): return (target.is_record() and can_cast_record(space, origin, target, casting)) else: if casting == 'equiv': return origin.elsize == target.elsize elif casting == 'safe': return origin.elsize <= target.elsize else: return True elif casting == 'same_kind': if can_cast_to(origin, target): return True if origin.kind in kind_ordering and target.kind in kind_ordering: return kind_ordering[origin.kind] <= kind_ordering[target.kind] return False elif casting == 'safe': return can_cast_to(origin, target) else: # 'equiv' return origin.num == target.num and origin.elsize == target.elsize def can_cast_record(space, origin, target, casting): if origin is target: return True if origin.fields is None or target.fields is None: return False if len(origin.fields) != len(target.fields): return False for name, (offset, orig_field) in origin.fields.iteritems(): if name not in target.fields: return False target_field = target.fields[name][1] if not can_cast_type(space, orig_field, target_field, casting): return False return True def can_cast_array(space, w_from, target, casting): # equivalent to PyArray_CanCastArrayTo origin = w_from.get_dtype() if w_from.is_scalar(): return can_cast_scalar( space, origin, w_from.get_scalar_value(), target, casting) else: return can_cast_type(space, origin, target, casting) def can_cast_scalar(space, from_type, value, target, casting): # equivalent to CNumPy's can_cast_scalar_to if from_type == target or casting == 'unsafe': return True if not from_type.is_number() or casting in ('no', 'equiv'): return can_cast_type(space, from_type, target, casting) if not from_type.is_native(): value = value.descr_byteswap(space) dtypenum, altnum = value.min_dtype() if target.is_unsigned(): dtypenum = altnum dtype = num2dtype(space, dtypenum) return can_cast_type(space, dtype, target, casting) def as_scalar(space, w_obj): dtype = scalar2dtype(space, w_obj) return dtype.coerce(space, w_obj) def min_scalar_type(space, w_a): w_array = convert_to_array(space, w_a) dtype = w_array.get_dtype() if w_array.is_scalar() and dtype.is_number(): num, alt_num = w_array.get_scalar_value().min_dtype() return num2dtype(space, num) else: return dtype def w_promote_types(space, w_type1, w_type2): dt1 = as_dtype(space, w_type1, allow_None=False) dt2 = as_dtype(space, w_type2, allow_None=False) return promote_types(space, dt1, dt2) def find_binop_result_dtype(space, dt1, dt2): if dt2 is None: return dt1 if dt1 is None: return dt2 return promote_types(space, dt1, dt2) def promote_types(space, dt1, dt2): """Return the smallest dtype to which both input dtypes can be safely cast""" # Equivalent to PyArray_PromoteTypes num = promotion_table[dt1.num][dt2.num] if num != -1: return num2dtype(space, num) # dt1.num should be <= dt2.num if dt1.num > dt2.num: dt1, dt2 = dt2, dt1 if dt2.is_str(): if dt1.is_str(): if dt1.elsize > dt2.elsize: return dt1 else: return dt2 else: # dt1 is numeric dt1_size = dt1.itemtype.strlen if dt1_size > dt2.elsize: return new_string_dtype(space, dt1_size) else: return dt2 elif dt2.is_unicode(): if dt1.is_unicode(): if dt1.elsize > dt2.elsize: return dt1 else: return dt2 elif dt1.is_str(): if dt2.elsize >= 4 * dt1.elsize: return dt2 else: return new_unicode_dtype(space, dt1.elsize) else: # dt1 is numeric dt1_size = dt1.itemtype.strlen if 4 * dt1_size > dt2.elsize: return new_unicode_dtype(space, dt1_size) else: return dt2 else: assert dt2.num == NPY.VOID if can_cast_type(space, dt1, dt2, casting='equiv'): return dt1 raise oefmt(space.w_TypeError, "invalid type promotion") def _promote_types_su(space, dt1, dt2, su1, su2): """Like promote_types(), but handles the small_unsigned flag as well""" if su1: if dt2.is_bool() or dt2.is_unsigned(): dt1 = dt1.as_unsigned(space) else: dt1 = dt1.as_signed(space) elif su2: if dt1.is_bool() or dt1.is_unsigned(): dt2 = dt2.as_unsigned(space) else: dt2 = dt2.as_signed(space) if dt1.elsize < dt2.elsize: su = su2 and (su1 or not dt1.is_signed()) elif dt1.elsize == dt2.elsize: su = su1 and su2 else: su = su1 and (su2 or not dt2.is_signed()) return promote_types(space, dt1, dt2), su def scalar2dtype(space, w_obj): from .boxes import W_GenericBox bool_dtype = get_dtype_cache(space).w_booldtype long_dtype = get_dtype_cache(space).w_longdtype int64_dtype = get_dtype_cache(space).w_int64dtype uint64_dtype = get_dtype_cache(space).w_uint64dtype complex_dtype = get_dtype_cache(space).w_complex128dtype float_dtype = get_dtype_cache(space).w_float64dtype object_dtype = get_dtype_cache(space).w_objectdtype if isinstance(w_obj, W_GenericBox): return w_obj.get_dtype(space) if space.isinstance_w(w_obj, space.w_bool): return bool_dtype elif space.isinstance_w(w_obj, space.w_int): try: space.int_w(w_obj) except OperationError as e: if e.match(space, space.w_OverflowError): if space.is_true(space.le(w_obj, space.newint(0))): return int64_dtype return uint64_dtype raise return int64_dtype elif space.isinstance_w(w_obj, space.w_float): return float_dtype elif space.isinstance_w(w_obj, space.w_complex): return complex_dtype elif space.isinstance_w(w_obj, space.w_bytes): return variable_dtype(space, 'S%d' % space.len_w(w_obj)) elif space.isinstance_w(w_obj, space.w_unicode): return new_unicode_dtype(space, space.len_w(w_obj)) return object_dtype @signature(ann.instance(W_Dtype), ann.instance(W_Dtype), returns=ann.bool()) def can_cast_to(dt1, dt2): """Return whether dtype `dt1` can be cast safely to `dt2`""" # equivalent to PyArray_CanCastTo from .casting import can_cast_itemtype result = can_cast_itemtype(dt1.itemtype, dt2.itemtype) if result: if dt1.num == NPY.STRING: if dt2.num == NPY.STRING: return dt1.elsize <= dt2.elsize elif dt2.num == NPY.UNICODE: return dt1.elsize * 4 <= dt2.elsize elif dt1.num == NPY.UNICODE and dt2.num == NPY.UNICODE: return dt1.elsize <= dt2.elsize elif dt2.num in (NPY.STRING, NPY.UNICODE): if dt2.num == NPY.STRING: char_size = 1 else: # NPY.UNICODE char_size = 4 if dt2.elsize == 0: return True if dt1.is_int(): return dt2.elsize >= dt1.itemtype.strlen * char_size return result @signature(ann.instance(BaseType), ann.instance(BaseType), returns=ann.bool()) def can_cast_itemtype(tp1, tp2): # equivalent to PyArray_CanCastSafely return casting_table[tp1.num][tp2.num] #_________________________ casting_table = [[False] * NPY.NTYPES for _ in range(NPY.NTYPES)] def enable_cast(type1, type2): casting_table[type1.num][type2.num] = True def _can_cast(type1, type2): """NOT_RPYTHON: operates on BaseType subclasses""" return casting_table[type1.num][type2.num] for tp in all_types: enable_cast(tp, tp) if tp.num != NPY.DATETIME: enable_cast(Bool, tp) enable_cast(tp, ObjectType) enable_cast(tp, VoidType) enable_cast(StringType, UnicodeType) #enable_cast(Bool, TimeDelta) for tp in number_types: enable_cast(tp, StringType) enable_cast(tp, UnicodeType) for tp1 in int_types: for tp2 in int_types: if tp1.signed: if tp2.signed and tp1.basesize() <= tp2.basesize(): enable_cast(tp1, tp2) else: if tp2.signed and tp1.basesize() < tp2.basesize(): enable_cast(tp1, tp2) elif not tp2.signed and tp1.basesize() <= tp2.basesize(): enable_cast(tp1, tp2) for tp1 in int_types: for tp2 in float_types + complex_types: size1 = tp1.basesize() size2 = tp2.basesize() if (size1 < 8 and size2 > size1) or (size1 >= 8 and size2 >= size1): enable_cast(tp1, tp2) for tp1 in float_types: for tp2 in float_types + complex_types: if tp1.basesize() <= tp2.basesize(): enable_cast(tp1, tp2) for tp1 in complex_types: for tp2 in complex_types: if tp1.basesize() <= tp2.basesize(): enable_cast(tp1, tp2) promotion_table = [[-1] * NPY.NTYPES for _ in range(NPY.NTYPES)] def promotes(tp1, tp2, tp3): if tp3 is None: num = -1 else: num = tp3.num promotion_table[tp1.num][tp2.num] = num for tp in all_types: promotes(tp, ObjectType, ObjectType) promotes(ObjectType, tp, ObjectType) for tp1 in [Bool] + number_types: for tp2 in [Bool] + number_types: if tp1 is tp2: promotes(tp1, tp1, tp1) elif _can_cast(tp1, tp2): promotes(tp1, tp2, tp2) elif _can_cast(tp2, tp1): promotes(tp1, tp2, tp1) else: # Brute-force search for the least upper bound result = None for tp3 in number_types: if _can_cast(tp1, tp3) and _can_cast(tp2, tp3): if result is None: result = tp3 elif _can_cast(tp3, result) and not _can_cast(result, tp3): result = tp3 promotes(tp1, tp2, result)