def _isna_array(values: ArrayLike, inf_as_na: bool = False):
"""
Return an array indicating which values of the input array are NaN / NA.
Parameters
----------
obj: ndarray or ExtensionArray
The input array whose elements are to be checked.
inf_as_na: bool
Whether or not to treat infinite values as NA.
Returns
-------
array-like
Array of boolean values denoting the NA status of each element.
"""
dtype = values.dtype
if is_extension_array_dtype(dtype):
if inf_as_na and is_categorical_dtype(dtype):
result = libmissing.isnaobj_old(values.to_numpy())
else:
result = values.isna()
elif is_string_dtype(dtype):
result = _isna_string_dtype(values, dtype, inf_as_na=inf_as_na)
elif needs_i8_conversion(dtype):
# this is the NaT pattern
result = values.view("i8") == iNaT
else:
if inf_as_na:
result = ~np.isfinite(values)
else:
result = np.isnan(values)
return result
def _isna_array(values: ArrayLike, inf_as_na: bool = False):
"""
Return an array indicating which values of the input array are NaN / NA.
Parameters
----------
obj: ndarray or ExtensionArray
The input array whose elements are to be checked.
inf_as_na: bool
Whether or not to treat infinite values as NA.
Returns
-------
array-like
Array of boolean values denoting the NA status of each element.
"""
dtype = values.dtype
if is_extension_array_dtype(dtype):
if inf_as_na and is_categorical_dtype(dtype):
# error: Item "ndarray" of "Union[ExtensionArray, ndarray]" has no attribute
# "to_numpy"
result = libmissing.isnaobj_old(
values.to_numpy() # type: ignore[union-attr]
)
else:
# error: Item "ndarray" of "Union[ExtensionArray, ndarray]" has no attribute
# "isna"
result = values.isna() # type: ignore[union-attr]
elif is_string_dtype(dtype):
# error: Argument 1 to "_isna_string_dtype" has incompatible type
# "ExtensionArray"; expected "ndarray"
# error: Argument 2 to "_isna_string_dtype" has incompatible type
# "ExtensionDtype"; expected "dtype[Any]"
result = _isna_string_dtype(
values,
dtype,
inf_as_na=inf_as_na # type: ignore[arg-type]
)
elif needs_i8_conversion(dtype):
# this is the NaT pattern
result = values.view("i8") == iNaT
else:
if inf_as_na:
# error: Argument 1 to "__call__" of "ufunc" has incompatible type
# "ExtensionArray"; expected "Union[Union[int, float, complex, str, bytes,
# generic], Sequence[Union[int, float, complex, str, bytes, generic]],
# Sequence[Sequence[Any]], _SupportsArray]"
result = ~np.isfinite(values) # type: ignore[arg-type]
else:
# error: Argument 1 to "__call__" of "ufunc" has incompatible type
# "ExtensionArray"; expected "Union[Union[int, float, complex, str, bytes,
# generic], Sequence[Union[int, float, complex, str, bytes, generic]],
# Sequence[Sequence[Any]], _SupportsArray]"
result = np.isnan(values) # type: ignore[arg-type]
return result
def extract_bool_array(mask: ArrayLike) -> npt.NDArray[np.bool_]:
"""
If we have a SparseArray or BooleanArray, convert it to ndarray[bool].
"""
if isinstance(mask, ExtensionArray):
# We could have BooleanArray, Sparse[bool], ...
# Except for BooleanArray, this is equivalent to just
# np.asarray(mask, dtype=bool)
mask = mask.to_numpy(dtype=bool, na_value=False)
mask = np.asarray(mask, dtype=bool)
return mask
def extract_bool_array(mask: ArrayLike) -> np.ndarray:
"""
If we have a SparseArray or BooleanArray, convert it to ndarray[bool].
"""
if isinstance(mask, ExtensionArray):
# We could have BooleanArray, Sparse[bool], ...
# Except for BooleanArray, this is equivalent to just
# np.asarray(mask, dtype=bool)
# error: Incompatible types in assignment (expression has type "ndarray",
# variable has type "ExtensionArray")
mask = mask.to_numpy(dtype=bool, na_value=False) # type: ignore[assignment]
# error: Incompatible types in assignment (expression has type "ndarray", variable
# has type "ExtensionArray")
mask = np.asarray(mask, dtype=bool) # type: ignore[assignment]
# error: Incompatible return value type (got "ExtensionArray", expected "ndarray")
return mask # type: ignore[return-value]
def ensure_int_or_float(arr: ArrayLike, copy: bool = False) -> np.ndarray:
"""
Ensure that an dtype array of some integer dtype
has an int64 dtype if possible.
If it's not possible, potentially because of overflow,
convert the array to float64 instead.
Parameters
----------
arr : array-like
The array whose data type we want to enforce.
copy: bool
Whether to copy the original array or reuse
it in place, if possible.
Returns
-------
out_arr : The input array cast as int64 if
possible without overflow.
Otherwise the input array cast to float64.
Notes
-----
If the array is explicitly of type uint64 the type
will remain unchanged.
"""
# TODO: GH27506 potential bug with ExtensionArrays
try:
# error: Unexpected keyword argument "casting" for "astype"
return arr.astype("int64", copy=copy,
casting="safe") # type: ignore[call-arg]
except TypeError:
pass
try:
# error: Unexpected keyword argument "casting" for "astype"
return arr.astype("uint64", copy=copy,
casting="safe") # type: ignore[call-arg]
except TypeError:
if is_extension_array_dtype(arr.dtype):
# pandas/core/dtypes/common.py:168: error: Item "ndarray" of
# "Union[ExtensionArray, ndarray]" has no attribute "to_numpy" [union-attr]
return arr.to_numpy( # type: ignore[union-attr]
dtype="float64", na_value=np.nan)
return arr.astype("float64", copy=copy)
def _isna_array(values: ArrayLike, inf_as_na: bool = False):
"""
Return an array indicating which values of the input array are NaN / NA.
Parameters
----------
obj: ndarray or ExtensionArray
The input array whose elements are to be checked.
inf_as_na: bool
Whether or not to treat infinite values as NA.
Returns
-------
array-like
Array of boolean values denoting the NA status of each element.
"""
dtype = values.dtype
if not isinstance(values, np.ndarray):
# i.e. ExtensionArray
if inf_as_na and is_categorical_dtype(dtype):
result = libmissing.isnaobj(values.to_numpy(), inf_as_na=inf_as_na)
else:
# error: Incompatible types in assignment (expression has type
# "Union[ndarray[Any, Any], ExtensionArraySupportsAnyAll]", variable has
# type "ndarray[Any, dtype[bool_]]")
result = values.isna() # type: ignore[assignment]
elif is_string_or_object_np_dtype(values.dtype):
result = _isna_string_dtype(values, inf_as_na=inf_as_na)
elif needs_i8_conversion(dtype):
# this is the NaT pattern
result = values.view("i8") == iNaT
else:
if inf_as_na:
result = ~np.isfinite(values)
else:
result = np.isnan(values)
return result
def ensure_int_or_float(arr: ArrayLike, copy: bool = False) -> np.array:
"""
Ensure that an dtype array of some integer dtype
has an int64 dtype if possible.
If it's not possible, potentially because of overflow,
convert the array to float64 instead.
Parameters
----------
arr : array-like
The array whose data type we want to enforce.
copy: bool
Whether to copy the original array or reuse
it in place, if possible.
Returns
-------
out_arr : The input array cast as int64 if
possible without overflow.
Otherwise the input array cast to float64.
Notes
-----
If the array is explicitly of type uint64 the type
will remain unchanged.
"""
# TODO: GH27506 potential bug with ExtensionArrays
try:
return arr.astype("int64", copy=copy, casting="safe") # type: ignore
except TypeError:
pass
try:
return arr.astype("uint64", copy=copy, casting="safe") # type: ignore
except TypeError:
if is_extension_array_dtype(arr.dtype):
return arr.to_numpy(dtype="float64", na_value=np.nan)
return arr.astype("float64", copy=copy)