1"""
2Methods used by Block.replace and related methods.
3"""
4
5from __future__ import annotations
6
7import operator
8import re
9from re import Pattern
10from typing import (
11 TYPE_CHECKING,
12 Any,
13)
14
15import numpy as np
16
17from pandas.core.dtypes.common import (
18 is_bool,
19 is_re,
20 is_re_compilable,
21)
22from pandas.core.dtypes.missing import isna
23
24if TYPE_CHECKING:
25 from pandas._typing import (
26 ArrayLike,
27 Scalar,
28 npt,
29 )
30
31
32def should_use_regex(regex: bool, to_replace: Any) -> bool:
33 """
34 Decide whether to treat `to_replace` as a regular expression.
35 """
36 if is_re(to_replace):
37 regex = True
38
39 regex = regex and is_re_compilable(to_replace)
40
41 # Don't use regex if the pattern is empty.
42 regex = regex and re.compile(to_replace).pattern != ""
43 return regex
44
45
46def compare_or_regex_search(
47 a: ArrayLike, b: Scalar | Pattern, regex: bool, mask: npt.NDArray[np.bool_]
48) -> ArrayLike:
49 """
50 Compare two array-like inputs of the same shape or two scalar values
51
52 Calls operator.eq or re.search, depending on regex argument. If regex is
53 True, perform an element-wise regex matching.
54
55 Parameters
56 ----------
57 a : array-like
58 b : scalar or regex pattern
59 regex : bool
60 mask : np.ndarray[bool]
61
62 Returns
63 -------
64 mask : array-like of bool
65 """
66 if isna(b):
67 return ~mask
68
69 def _check_comparison_types(
70 result: ArrayLike | bool, a: ArrayLike, b: Scalar | Pattern
71 ) -> None:
72 """
73 Raises an error if the two arrays (a,b) cannot be compared.
74 Otherwise, returns the comparison result as expected.
75 """
76 if is_bool(result) and isinstance(a, np.ndarray):
77 type_names = [type(a).__name__, type(b).__name__]
78
79 type_names[0] = f"ndarray(dtype={a.dtype})"
80
81 raise TypeError(
82 f"Cannot compare types {type_names[0]!r} and {type_names[1]!r}"
83 )
84
85 if not regex or not should_use_regex(regex, b):
86 # TODO: should use missing.mask_missing?
87 op = lambda x: operator.eq(x, b)
88 else:
89 op = np.vectorize(
90 lambda x: (
91 bool(re.search(b, x))
92 if isinstance(x, str) and isinstance(b, (str, Pattern))
93 else False
94 ),
95 otypes=[bool],
96 )
97
98 # GH#32621 use mask to avoid comparing to NAs
99 if isinstance(a, np.ndarray) and mask is not None:
100 a = a[mask]
101 result = op(a)
102
103 if isinstance(result, np.ndarray):
104 # The shape of the mask can differ to that of the result
105 # since we may compare only a subset of a's or b's elements
106 tmp = np.zeros(mask.shape, dtype=np.bool_)
107 np.place(tmp, mask, result)
108 result = tmp
109 else:
110 result = op(a)
111
112 _check_comparison_types(result, a, b)
113 return result
114
115
116def replace_regex(
117 values: ArrayLike, rx: re.Pattern, value, mask: npt.NDArray[np.bool_] | None
118) -> None:
119 """
120 Parameters
121 ----------
122 values : ArrayLike
123 Object dtype.
124 rx : re.Pattern
125 value : Any
126 mask : np.ndarray[bool], optional
127
128 Notes
129 -----
130 Alters values in-place.
131 """
132
133 # deal with replacing values with objects (strings) that match but
134 # whose replacement is not a string (numeric, nan, object)
135 if isna(value) or not isinstance(value, str):
136
137 def re_replacer(s):
138 if is_re(rx) and isinstance(s, str):
139 return value if rx.search(s) is not None else s
140 else:
141 return s
142
143 else:
144 # value is guaranteed to be a string here, s can be either a string
145 # or null if it's null it gets returned
146 def re_replacer(s):
147 if is_re(rx) and isinstance(s, str):
148 return rx.sub(value, s)
149 else:
150 return s
151
152 f = np.vectorize(re_replacer, otypes=[np.object_])
153
154 if mask is None:
155 values[:] = f(values)
156 else:
157 if values.ndim != mask.ndim:
158 mask = np.broadcast_to(mask, values.shape)
159 values[mask] = f(values[mask])