add regex support to _replace_coerce

Author: peterpanmj

pandas/core/internals.py

@@ -1690,7 +1690,8 @@ def _nanpercentile(values, q, axis, **kw):
                               placement=np.arange(len(result)),
                               ndim=ndim)
 
-    def _replace_coerce(self, mask=None, dst=None, convert=False):
+    def _replace_coerce(self, mask=None, src=None, dst=None, inplace=True,
+                        convert=False, regex=False, mgr=None):
         """
         Replace value corresponding to the given boolean array with another
         value.
@@ -1699,18 +1700,32 @@ def _replace_coerce(self, mask=None, dst=None, convert=False):
         ----------
         mask : array_like of bool
             The mask of values to replace.
+        src : object
+            The value to replace. It is ignored if regex is False.
         dst : object
             The value to be replaced with.
         convert : bool
-            It is used in ObjectBlocks. It is here for API compatibility.
+            If true, try to coerce any object types to better types.
+        regex : bool
+            If true, search for element matching with the pattern in src.
+            Masked element is ignored.
+        mgr : BlockPlacement, optional
 
         Returns
         -------
         A new block if there is anything to replace or the original block.
         """
+
         if mask.any():
-            self = self.coerce_to_target_dtype(dst)
-            return self.putmask(mask, dst, inplace=True)
+            if not regex:
+                self = self.coerce_to_target_dtype(dst)
+                return self.putmask(mask, dst, inplace=inplace)
+            else:
+                return self._replace_single(src, dst, inplace=inplace,
+                                            regex=regex,
+                                            convert=convert,
+                                            mask=mask,
+                                            mgr=mgr)
         return self
 
 
@@ -2488,7 +2503,7 @@ def replace(self, to_replace, value, inplace=False, filter=None,
                                     regex=regex, mgr=mgr)
 
     def _replace_single(self, to_replace, value, inplace=False, filter=None,
-                        regex=False, convert=True, mgr=None):
+                        regex=False, convert=True, mgr=None, mask=None):
 
         inplace = validate_bool_kwarg(inplace, 'inplace')
 
@@ -2555,15 +2570,19 @@ def re_replacer(s):
         else:
             filt = self.mgr_locs.isin(filter).nonzero()[0]
 
-        new_values[filt] = f(new_values[filt])
+        if mask is None:
+            new_values[filt] = f(new_values[filt])
+        else:
+            new_values[filt][mask] = f(new_values[filt][mask])
 
         # convert
         block = self.make_block(new_values)
         if convert:
             block = block.convert(by_item=True, numeric=False)
         return block
 
-    def _replace_coerce(self, mask=None, dst=None, convert=False):
+    def _replace_coerce(self, mask=None, src=None, dst=None, inplace=True,
+                        convert=False, regex=False, mgr=None):
         """
         Replace value corresponding to the given boolean array with another
         value.
@@ -2572,17 +2591,29 @@ def _replace_coerce(self, mask=None, dst=None, convert=False):
         ----------
         mask : array_like of bool
             The mask of values to replace.
+        src : object
+            The value to replace. It is ignored if regex is False.
         dst : object
             The value to be replaced with.
         convert : bool
             If true, try to coerce any object types to better types.
+        regex : bool
+            If true, search for element matching with the pattern in src.
+            Masked element is ignored.
+        mgr : BlockPlacement, optional
 
         Returns
         -------
         A new block if there is anything to replace or the original block.
         """
         if mask.any():
-            block = super(ObjectBlock, self)._replace_coerce(mask, dst)
+            block = super(ObjectBlock, self)._replace_coerce(mask=mask,
+                                                             src=src,
+                                                             dst=dst,
+                                                             inplace=inplace,
+                                                             convert=convert,
+                                                             regex=regex,
+                                                             mgr=mgr)
             if convert:
                 block = [b.convert(by_item=True, numeric=False, copy=True)
                          for b in block]
@@ -3779,17 +3810,17 @@ def replace_list(self, src_list, dest_list, inplace=False, regex=False,
         # figure out our mask a-priori to avoid repeated replacements
         values = self.as_array()
 
-        # only support equality comparision, regex comparision support
-        # is needed in the future
-        def comp(s):
+        def comp(s, reg=False):
             if isna(s):
                 return isna(values)
             if hasattr(s, 'asm8'):
                 return _maybe_compare(maybe_convert_objects(values),
-                                      getattr(s, 'asm8'), operator.eq)
-            return _maybe_compare(values, s, operator.eq)
+                                      getattr(s, 'asm8'), reg)
+            if reg and is_re_compilable(s):
+                return _maybe_compare(values, s, reg)
+            return _maybe_compare(values, s, reg)
 
-        masks = [comp(s) for i, s in enumerate(src_list)]
+        masks = [comp(s, regex) for i, s in enumerate(src_list)]
 
         result_blocks = []
         src_len = len(src_list) - 1
@@ -3801,39 +3832,16 @@ def comp(s):
             for i, (s, d) in enumerate(zip(src_list, dest_list)):
                 new_rb = []
                 for b in rb:
-                    # regular expression support needs to be improved.
-                    # If the replacement for the previous pattern
-                    # matches the next pattern, the value will be replaced
-                    # again with a different value from dest_list
-                    # i.e. when values is ['a', 'b']
-                    # src_list: [r'a*', r'b*'], dest_list: ['b', 'a']
-                    # result will be ['b', b'] after searching for pattern r'a'
-                    # and then changed to ['a', 'a'] for pattern r'b*'
-                    if regex:
-                        if is_object_dtype(b.dtype):
-                            convert = i == src_len
-                            result = b.replace(s, d, inplace=inplace,
-                                               regex=regex,
-                                               mgr=mgr, convert=convert)
-                            new_rb = _extend_blocks(result, new_rb)
-                        else:
-                            # get our mask for this element, sized to this
-                            # particular block
-                            m = masks[i][b.mgr_locs.indexer]
-                            if m.any():
-                                b = b.coerce_to_target_dtype(d)
-                                new_rb.extend(b.putmask(m, d, inplace=True))
-                            else:
-                                new_rb.append(b)
+                    m = masks[i][b.mgr_locs.indexer]
+                    convert = i == src_len
+                    result = b._replace_coerce(mask=m, src=s, dst=d,
+                                               inplace=inplace,
+                                               convert=convert, regex=regex,
+                                               mgr=mgr)
+                    if m.any():
+                        new_rb = _extend_blocks(result, new_rb)
                     else:
-                        m = masks[i][b.mgr_locs.indexer]
-                        convert = i == src_len
-                        result = b._replace_coerce(mask=m, dst=d,
-                                                   convert=convert)
-                        if m.any():
-                            new_rb = _extend_blocks(result, new_rb)
-                        else:
-                            new_rb.append(b)
+                        new_rb.append(b)
                 rb = new_rb
             result_blocks.extend(rb)
 
@@ -5207,7 +5215,12 @@ def _vstack(to_stack, dtype):
         return np.vstack(to_stack)
 
 
-def _maybe_compare(a, b, op):
+def _maybe_compare(a, b, regex=False):
+    if not regex:
+        op = lambda x: operator.eq(x, b)
+    else:
+        op = np.vectorize(lambda x: bool(re.match(b, x)) if isinstance(x, str)
+                          else False)
 
     is_a_array = isinstance(a, np.ndarray)
     is_b_array = isinstance(b, np.ndarray)
@@ -5219,9 +5232,8 @@ def _maybe_compare(a, b, op):
     # numpy deprecation warning if comparing numeric vs string-like
     elif is_numeric_v_string_like(a, b):
         result = False
-
     else:
-        result = op(a, b)
+        result = op(a)
 
     if is_scalar(result) and (is_a_array or is_b_array):
         type_names = [type(a).__name__, type(b).__name__]