def putmask_smart(values: np.ndarray, mask: np.ndarray, new) -> np.ndarray: """ Return a new ndarray, try to preserve dtype if possible. Parameters ---------- values : np.ndarray `values`, updated in-place. mask : np.ndarray[bool] Applies to both sides (array like). new : `new values` either scalar or an array like aligned with `values` Returns ------- values : ndarray with updated values this *may* be a copy of the original See Also -------- ndarray.putmask """ # we cannot use np.asarray() here as we cannot have conversions # that numpy does when numeric are mixed with strings # n should be the length of the mask or a scalar here if not is_list_like(new): new = np.repeat(new, len(mask)) # see if we are only masking values that if putted # will work in the current dtype try: nn = new[mask] except TypeError: # TypeError: only integer scalar arrays can be converted to a scalar index pass else: # make sure that we have a nullable type if we have nulls if not isna_compat(values, nn[0]): pass elif not (is_float_dtype(nn.dtype) or is_integer_dtype(nn.dtype)): # only compare integers/floats pass elif not (is_float_dtype(values.dtype) or is_integer_dtype(values.dtype)): # only compare integers/floats pass else: # we ignore ComplexWarning here with warnings.catch_warnings(record=True): warnings.simplefilter("ignore", np.ComplexWarning) nn_at = nn.astype(values.dtype) comp = nn == nn_at if is_list_like(comp) and comp.all(): nv = values.copy() nv[mask] = nn_at return nv new = np.asarray(new) if values.dtype.kind == new.dtype.kind: # preserves dtype if possible return _putmask_preserve(values, new, mask) # change the dtype if needed dtype, _ = maybe_promote(new.dtype) values = values.astype(dtype) return _putmask_preserve(values, new, mask)
def putmask_smart(values: np.ndarray, mask: np.ndarray, new) -> np.ndarray: """ Return a new ndarray, try to preserve dtype if possible. Parameters ---------- values : np.ndarray `values`, updated in-place. mask : np.ndarray[bool] Applies to both sides (array like). new : `new values` either scalar or an array like aligned with `values` Returns ------- values : ndarray with updated values this *may* be a copy of the original See Also -------- ndarray.putmask """ # we cannot use np.asarray() here as we cannot have conversions # that numpy does when numeric are mixed with strings # n should be the length of the mask or a scalar here if not is_list_like(new): new = np.broadcast_to(new, mask.shape) # see if we are only masking values that if putted # will work in the current dtype try: nn = new[mask] except TypeError: # TypeError: only integer scalar arrays can be converted to a scalar index pass else: # make sure that we have a nullable type if we have nulls if not isna_compat(values, nn[0]): pass elif not (is_float_dtype(nn.dtype) or is_integer_dtype(nn.dtype)): # only compare integers/floats pass elif not (is_float_dtype(values.dtype) or is_integer_dtype(values.dtype)): # only compare integers/floats pass else: # we ignore ComplexWarning here with warnings.catch_warnings(record=True): warnings.simplefilter("ignore", np.ComplexWarning) nn_at = nn.astype(values.dtype) comp = nn == nn_at if is_list_like(comp) and comp.all(): nv = values.copy() nv[mask] = nn_at return nv new = np.asarray(new) if values.dtype.kind == new.dtype.kind: # preserves dtype if possible return _putmask_preserve(values, new, mask) dtype = find_common_type([values.dtype, new.dtype]) # error: Argument 1 to "astype" of "_ArrayOrScalarCommon" has incompatible type # "Union[dtype[Any], ExtensionDtype]"; expected "Union[dtype[Any], None, type, # _SupportsDType, str, Union[Tuple[Any, int], Tuple[Any, Union[int, Sequence[int]]], # List[Any], _DTypeDict, Tuple[Any, Any]]]" values = values.astype(dtype) # type: ignore[arg-type] return _putmask_preserve(values, new, mask)