Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/pandas/core/array_algos/putmask.py: 28%

1""" 

2EA-compatible analogue to np.putmask 

3""" 

4from __future__ import annotations 

5 

6from typing import ( 

7 TYPE_CHECKING, 

8 Any, 

9) 

10 

11import numpy as np 

12 

13from pandas._libs import lib 

14 

15from pandas.core.dtypes.cast import infer_dtype_from 

16from pandas.core.dtypes.common import is_list_like 

17 

18from pandas.core.arrays import ExtensionArray 

19 

20if TYPE_CHECKING: 

21 from pandas._typing import ( 

22 ArrayLike, 

23 npt, 

24 ) 

25 

26 from pandas import MultiIndex 

27 

28 

29def putmask_inplace(values: ArrayLike, mask: npt.NDArray[np.bool_], value: Any) -> None: 

30 """ 

31 ExtensionArray-compatible implementation of np.putmask. The main 

32 difference is we do not handle repeating or truncating like numpy. 

33 

34 Parameters 

35 ---------- 

36 values: np.ndarray or ExtensionArray 

37 mask : np.ndarray[bool] 

38 We assume extract_bool_array has already been called. 

39 value : Any 

40 """ 

41 

42 if ( 

43 not isinstance(values, np.ndarray) 

44 or (values.dtype == object and not lib.is_scalar(value)) 

45 # GH#43424: np.putmask raises TypeError if we cannot cast between types with 

46 # rule = "safe", a stricter guarantee we may not have here 

47 or ( 

48 isinstance(value, np.ndarray) and not np.can_cast(value.dtype, values.dtype) 

49 ) 

50 ): 

51 # GH#19266 using np.putmask gives unexpected results with listlike value 

52 # along with object dtype 

53 if is_list_like(value) and len(value) == len(values): 

54 values[mask] = value[mask] 

55 else: 

56 values[mask] = value 

57 else: 

58 # GH#37833 np.putmask is more performant than __setitem__ 

59 np.putmask(values, mask, value) 

60 

61 

62def putmask_without_repeat( 

63 values: np.ndarray, mask: npt.NDArray[np.bool_], new: Any 

64) -> None: 

65 """ 

66 np.putmask will truncate or repeat if `new` is a listlike with 

67 len(new) != len(values). We require an exact match. 

68 

69 Parameters 

70 ---------- 

71 values : np.ndarray 

72 mask : np.ndarray[bool] 

73 new : Any 

74 """ 

75 if getattr(new, "ndim", 0) >= 1: 

76 new = new.astype(values.dtype, copy=False) 

77 

78 # TODO: this prob needs some better checking for 2D cases 

79 nlocs = mask.sum() 

80 if nlocs > 0 and is_list_like(new) and getattr(new, "ndim", 1) == 1: 

81 shape = np.shape(new) 

82 # np.shape compat for if setitem_datetimelike_compat 

83 # changed arraylike to list e.g. test_where_dt64_2d 

84 if nlocs == shape[-1]: 

85 # GH#30567 

86 # If length of ``new`` is less than the length of ``values``, 

87 # `np.putmask` would first repeat the ``new`` array and then 

88 # assign the masked values hence produces incorrect result. 

89 # `np.place` on the other hand uses the ``new`` values at it is 

90 # to place in the masked locations of ``values`` 

91 np.place(values, mask, new) 

92 # i.e. values[mask] = new 

93 elif mask.shape[-1] == shape[-1] or shape[-1] == 1: 

94 np.putmask(values, mask, new) 

95 else: 

96 raise ValueError("cannot assign mismatch length to masked array") 

97 else: 

98 np.putmask(values, mask, new) 

99 

100 

101def validate_putmask( 

102 values: ArrayLike | MultiIndex, mask: np.ndarray 

103) -> tuple[npt.NDArray[np.bool_], bool]: 

104 """ 

105 Validate mask and check if this putmask operation is a no-op. 

106 """ 

107 mask = extract_bool_array(mask) 

108 if mask.shape != values.shape: 

109 raise ValueError("putmask: mask and data must be the same size") 

110 

111 noop = not mask.any() 

112 return mask, noop 

113 

114 

115def extract_bool_array(mask: ArrayLike) -> npt.NDArray[np.bool_]: 

116 """ 

117 If we have a SparseArray or BooleanArray, convert it to ndarray[bool]. 

118 """ 

119 if isinstance(mask, ExtensionArray): 

120 # We could have BooleanArray, Sparse[bool], ... 

121 # Except for BooleanArray, this is equivalent to just 

122 # np.asarray(mask, dtype=bool) 

123 mask = mask.to_numpy(dtype=bool, na_value=False) 

124 

125 mask = np.asarray(mask, dtype=bool) 

126 return mask 

127 

128 

129def setitem_datetimelike_compat(values: np.ndarray, num_set: int, other): 

130 """ 

131 Parameters 

132 ---------- 

133 values : np.ndarray 

134 num_set : int 

135 For putmask, this is mask.sum() 

136 other : Any 

137 """ 

138 if values.dtype == object: 

139 dtype, _ = infer_dtype_from(other) 

140 

141 if lib.is_np_dtype(dtype, "mM"): 

142 # https://github.com/numpy/numpy/issues/12550 

143 # timedelta64 will incorrectly cast to int 

144 if not is_list_like(other): 

145 other = [other] * num_set 

146 else: 

147 other = list(other) 

148 

149 return other