20 lines
570 B
Python
20 lines
570 B
Python
|
import numpy as np
|
||
|
|
|
||
|
|
from . import is_scalar_nan
|
||
|
|
from .fixes import _object_dtype_isnan
|
||
|
|
|
||
|
|
|
||
|
|
def _get_mask(X, value_to_mask):
|
||
|
|
"""Compute the boolean mask X == value_to_mask."""
|
||
|
|
if is_scalar_nan(value_to_mask):
|
||
|
|
if X.dtype.kind == "f":
|
||
|
|
return np.isnan(X)
|
||
|
|
elif X.dtype.kind in ("i", "u"):
|
||
|
|
# can't have NaNs in integer array.
|
||
|
|
return np.zeros(X.shape, dtype=bool)
|
||
|
|
else:
|
||
|
|
# np.isnan does not work on object dtypes.
|
||
|
|
return _object_dtype_isnan(X)
|
||
|
|
else:
|
||
|
|
return X == value_to_mask
|