Coverage for /var/srv/projects/api.amasfac.comuna18.com/tmp/venv/lib/python3.9/site-packages/pandas/core/array_algos/putmask.py: 19%

1""" 

2EA-compatible analogue to np.putmask 

3""" 

4from __future__ import annotations 

5 

6from typing import Any 

7 

8import numpy as np 

9 

10from pandas._libs import lib 

11from pandas._typing import ( 

12 ArrayLike, 

13 npt, 

14) 

15from pandas.compat import np_version_under1p21 

16 

17from pandas.core.dtypes.cast import infer_dtype_from 

18from pandas.core.dtypes.common import is_list_like 

19 

20from pandas.core.arrays import ExtensionArray 

21 

22 

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

24 """ 

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

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

27 

28 Parameters 

29 ---------- 

30 values: np.ndarray or ExtensionArray 

31 mask : np.ndarray[bool] 

32 We assume extract_bool_array has already been called. 

33 value : Any 

34 """ 

35 

36 if ( 

37 not isinstance(values, np.ndarray) 

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

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

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

41 or ( 

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

43 ) 

44 ): 

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

46 # along with object dtype 

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

48 values[mask] = value[mask] 

49 else: 

50 values[mask] = value 

51 else: 

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

53 np.putmask(values, mask, value) 

54 

55 

56def putmask_without_repeat( 

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

58) -> None: 

59 """ 

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

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

62 

63 Parameters 

64 ---------- 

65 values : np.ndarray 

66 mask : np.ndarray[bool] 

67 new : Any 

68 """ 

69 if np_version_under1p21: 

70 new = setitem_datetimelike_compat(values, mask.sum(), new) 

71 

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

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

74 

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

76 nlocs = mask.sum() 

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

78 shape = np.shape(new) 

79 # np.shape compat for if setitem_datetimelike_compat 

80 # changed arraylike to list e.g. test_where_dt64_2d 

81 if nlocs == shape[-1]: 

82 # GH#30567 

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

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

85 # assign the masked values hence produces incorrect result. 

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

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

88 np.place(values, mask, new) 

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

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

91 np.putmask(values, mask, new) 

92 else: 

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

94 else: 

95 np.putmask(values, mask, new) 

96 

97 

98def validate_putmask( 

99 values: ArrayLike, mask: np.ndarray 

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

101 """ 

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

103 """ 

104 mask = extract_bool_array(mask) 

105 if mask.shape != values.shape: 

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

107 

108 noop = not mask.any() 

109 return mask, noop 

110 

111 

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

113 """ 

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

115 """ 

116 if isinstance(mask, ExtensionArray): 

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

118 # Except for BooleanArray, this is equivalent to just 

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

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

121 

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

123 return mask 

124 

125 

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

127 """ 

128 Parameters 

129 ---------- 

130 values : np.ndarray 

131 num_set : int 

132 For putmask, this is mask.sum() 

133 other : Any 

134 """ 

135 if values.dtype == object: 

136 dtype, _ = infer_dtype_from(other, pandas_dtype=True) 

137 

138 if isinstance(dtype, np.dtype) and dtype.kind in ["m", "M"]: 

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

140 # timedelta64 will incorrectly cast to int 

141 if not is_list_like(other): 

142 other = [other] * num_set 

143 else: 

144 other = list(other) 

145 

146 return other