Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/array_algos/putmask.py: 29%

1""" 

2EA-compatible analogue to np.putmask 

3""" 

4 

5from __future__ import annotations 

6 

7from typing import ( 

8 TYPE_CHECKING, 

9 Any, 

10) 

11 

12import numpy as np 

13 

14from pandas._libs import lib 

15 

16from pandas.core.dtypes.cast import infer_dtype_from 

17from pandas.core.dtypes.common import is_list_like 

18 

19from pandas.core.arrays import ExtensionArray 

20 

21if TYPE_CHECKING: 

22 from pandas._typing import ( 

23 ArrayLike, 

24 npt, 

25 ) 

26 

27 from pandas import MultiIndex 

28 

29 

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

31 """ 

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

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

34 

35 Parameters 

36 ---------- 

37 values: np.ndarray or ExtensionArray 

38 mask : np.ndarray[bool] 

39 We assume extract_bool_array has already been called. 

40 value : Any 

41 """ 

42 

43 if ( 

44 not isinstance(values, np.ndarray) 

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

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

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

48 or ( 

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

50 ) 

51 ): 

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

53 # along with object dtype 

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

55 values[mask] = value[mask] 

56 else: 

57 values[mask] = value 

58 else: 

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

60 np.putmask(values, mask, value) 

61 

62 

63def putmask_without_repeat( 

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

65) -> None: 

66 """ 

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

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

69 

70 Parameters 

71 ---------- 

72 values : np.ndarray 

73 mask : np.ndarray[bool] 

74 new : Any 

75 """ 

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

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

78 

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

80 nlocs = mask.sum() 

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

82 shape = np.shape(new) 

83 # np.shape compat for if setitem_datetimelike_compat 

84 # changed arraylike to list e.g. test_where_dt64_2d 

85 if nlocs == shape[-1]: 

86 # GH#30567 

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

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

89 # assign the masked values hence produces incorrect result. 

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

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

92 np.place(values, mask, new) 

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

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

95 np.putmask(values, mask, new) 

96 else: 

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

98 else: 

99 np.putmask(values, mask, new) 

100 

101 

102def validate_putmask( 

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

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

105 """ 

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

107 """ 

108 mask = extract_bool_array(mask) 

109 if mask.shape != values.shape: 

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

111 

112 noop = not mask.any() 

113 return mask, noop 

114 

115 

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

117 """ 

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

119 """ 

120 if isinstance(mask, ExtensionArray): 

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

122 # Except for BooleanArray, this is equivalent to just 

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

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

125 

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

127 return mask 

128 

129 

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

131 """ 

132 Parameters 

133 ---------- 

134 values : np.ndarray 

135 num_set : int 

136 For putmask, this is mask.sum() 

137 other : Any 

138 """ 

139 if values.dtype == object: 

140 dtype, _ = infer_dtype_from(other) 

141 

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

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

144 # timedelta64 will incorrectly cast to int 

145 if not is_list_like(other): 

146 other = [other] * num_set 

147 else: 

148 other = list(other) 

149 

150 return other