REF/TYP: consistent return type for Block.replace (#37010)

Author: jbrockmendel

pandas/core/internals/blocks.py

@@ -717,7 +717,7 @@ def replace(
         inplace: bool = False,
         regex: bool = False,
         convert: bool = True,
-    ):
+    ) -> List["Block"]:
         """
         replace the to_replace value with value, possible to create new
         blocks here this is just a call to putmask. regex is not used here.
@@ -808,9 +808,11 @@ def replace(
             )
         return blocks
 
-    def _replace_single(self, *args, **kwargs):
+    def _replace_single(
+        self, to_replace, value, inplace=False, regex=False, convert=True, mask=None
+    ) -> List["Block"]:
         """ no-op on a non-ObjectBlock """
-        return self if kwargs["inplace"] else self.copy()
+        return [self] if inplace else [self.copy()]
 
     def _replace_list(
         self,
@@ -852,16 +854,13 @@ def comp(s: Scalar, mask: np.ndarray, regex: bool = False) -> np.ndarray:
                     to_replace=src,
                     value=dest,
                     inplace=inplace,
-                    convert=convert,
                     regex=regex,
                 )
-                if m.any() or convert:
-                    if isinstance(result, list):
-                        new_rb.extend(result)
-                    else:
-                        new_rb.append(result)
-                else:
-                    new_rb.append(blk)
+                if convert and blk.is_object:
+                    result = extend_blocks(
+                        [b.convert(numeric=False, copy=True) for b in result]
+                    )
+                new_rb.extend(result)
             rb = new_rb
         return rb
 
@@ -1559,9 +1558,8 @@ def _replace_coerce(
         value,
         inplace: bool = True,
         regex: bool = False,
-        convert: bool = False,
         mask=None,
-    ):
+    ) -> List["Block"]:
         """
         Replace value corresponding to the given boolean array with another
         value.
@@ -1576,14 +1574,12 @@ def _replace_coerce(
             Perform inplace modification.
         regex : bool, default False
             If true, perform regular expression substitution.
-        convert : bool, default True
-            If true, try to coerce any object types to better types.
         mask : array-like of bool, optional
             True indicate corresponding element is ignored.
 
         Returns
         -------
-        A new block if there is anything to replace or the original block.
+        List[Block]
         """
         if mask.any():
             if not regex:
@@ -1595,10 +1591,10 @@ def _replace_coerce(
                     value,
                     inplace=inplace,
                     regex=regex,
-                    convert=convert,
+                    convert=False,
                     mask=mask,
                 )
-        return self
+        return [self]
 
 
 class ExtensionBlock(Block):
@@ -2479,14 +2475,16 @@ def _maybe_downcast(self, blocks: List["Block"], downcast=None) -> List["Block"]
     def _can_hold_element(self, element: Any) -> bool:
         return True
 
-    def replace(self, to_replace, value, inplace=False, regex=False, convert=True):
+    def replace(
+        self, to_replace, value, inplace=False, regex=False, convert=True
+    ) -> List["Block"]:
         to_rep_is_list = is_list_like(to_replace)
         value_is_list = is_list_like(value)
         both_lists = to_rep_is_list and value_is_list
         either_list = to_rep_is_list or value_is_list
 
-        result_blocks = []
-        blocks = [self]
+        result_blocks: List["Block"] = []
+        blocks: List["Block"] = [self]
 
         if not either_list and is_re(to_replace):
             return self._replace_single(
@@ -2503,7 +2501,7 @@ def replace(self, to_replace, value, inplace=False, regex=False, convert=True):
                     result = b._replace_single(
                         to_rep, v, inplace=inplace, regex=regex, convert=convert
                     )
-                    result_blocks = extend_blocks(result, result_blocks)
+                    result_blocks.extend(result)
                 blocks = result_blocks
             return result_blocks
 
@@ -2514,7 +2512,7 @@ def replace(self, to_replace, value, inplace=False, regex=False, convert=True):
                     result = b._replace_single(
                         to_rep, value, inplace=inplace, regex=regex, convert=convert
                     )
-                    result_blocks = extend_blocks(result, result_blocks)
+                    result_blocks.extend(result)
                 blocks = result_blocks
             return result_blocks
 
@@ -2524,7 +2522,7 @@ def replace(self, to_replace, value, inplace=False, regex=False, convert=True):
 
     def _replace_single(
         self, to_replace, value, inplace=False, regex=False, convert=True, mask=None
-    ):
+    ) -> List["Block"]:
         """
         Replace elements by the given value.
 
@@ -2545,7 +2543,7 @@ def _replace_single(
 
         Returns
         -------
-        a new block, the result after replacing
+        List[Block]
         """
         inplace = validate_bool_kwarg(inplace, "inplace")
 
@@ -2619,48 +2617,6 @@ def re_replacer(s):
             nbs = [block]
         return nbs
 
-    def _replace_coerce(
-        self, to_replace, value, inplace=True, regex=False, convert=False, mask=None
-    ):
-        """
-        Replace value corresponding to the given boolean array with another
-        value.
-
-        Parameters
-        ----------
-        to_replace : object or pattern
-            Scalar to replace or regular expression to match.
-        value : object
-            Replacement object.
-        inplace : bool, default False
-            Perform inplace modification.
-        regex : bool, default False
-            If true, perform regular expression substitution.
-        convert : bool, default True
-            If true, try to coerce any object types to better types.
-        mask : array-like of bool, optional
-            True indicate corresponding element is ignored.
-
-        Returns
-        -------
-        A new block if there is anything to replace or the original block.
-        """
-        if mask.any():
-            nbs = super()._replace_coerce(
-                to_replace=to_replace,
-                value=value,
-                inplace=inplace,
-                regex=regex,
-                convert=convert,
-                mask=mask,
-            )
-            if convert:
-                nbs = extend_blocks([b.convert(numeric=False, copy=True) for b in nbs])
-            return nbs
-        if convert:
-            return self.convert(numeric=False, copy=True)
-        return [self]
-
 
 class CategoricalBlock(ExtensionBlock):
     __slots__ = ()
@@ -2672,12 +2628,12 @@ def replace(
         inplace: bool = False,
         regex: bool = False,
         convert: bool = True,
-    ):
+    ) -> List["Block"]:
         inplace = validate_bool_kwarg(inplace, "inplace")
         result = self if inplace else self.copy()
 
         result.values.replace(to_replace, value, inplace=True)
-        return result
+        return [result]
 
 
 # -----------------------------------------------------------------