REF: de-duplicate block_shape vs safe_reshape (#39892)

Author: jbrockmendel

pandas/core/internals/__init__.py

@@ -11,7 +11,6 @@
     ObjectBlock,
     TimeDeltaBlock,
     make_block,
-    safe_reshape,
 )
 from pandas.core.internals.concat import concatenate_block_managers
 from pandas.core.internals.managers import (
@@ -31,7 +30,6 @@
     "FloatBlock",
     "ObjectBlock",
     "TimeDeltaBlock",
-    "safe_reshape",
     "make_block",
     "DataManager",
     "ArrayManager",

pandas/core/internals/blocks.py

@@ -308,7 +308,7 @@ def make_block(self, values, placement=None) -> Block:
         if placement is None:
             placement = self.mgr_locs
         if self.is_extension:
-            values = _block_shape(values, ndim=self.ndim)
+            values = ensure_block_shape(values, ndim=self.ndim)
 
         return make_block(values, placement=placement, ndim=self.ndim)
 
@@ -533,7 +533,7 @@ def make_a_block(nv, ref_loc):
             else:
                 # Put back the dimension that was taken from it and make
                 # a block out of the result.
-                nv = _block_shape(nv, ndim=self.ndim)
+                nv = ensure_block_shape(nv, ndim=self.ndim)
                 block = self.make_block(values=nv, placement=ref_loc)
             return block
 
@@ -1569,7 +1569,9 @@ def putmask(self, mask, new) -> List[Block]:
         if isinstance(new, (np.ndarray, ExtensionArray)) and len(new) == len(mask):
             new = new[mask]
 
-        mask = safe_reshape(mask, new_values.shape)
+        if mask.ndim == new_values.ndim + 1:
+            # TODO(EA2D): unnecessary with 2D EAs
+            mask = mask.reshape(new_values.shape)
 
         new_values[mask] = new
         return [self.make_block(values=new_values)]
@@ -2426,36 +2428,15 @@ def extend_blocks(result, blocks=None) -> List[Block]:
     return blocks
 
 
-def _block_shape(values: ArrayLike, ndim: int = 1) -> ArrayLike:
-    """ guarantee the shape of the values to be at least 1 d """
+def ensure_block_shape(values: ArrayLike, ndim: int = 1) -> ArrayLike:
+    """
+    Reshape if possible to have values.ndim == ndim.
+    """
     if values.ndim < ndim:
-        shape = values.shape
         if not is_extension_array_dtype(values.dtype):
             # TODO(EA2D): https://github.com/pandas-dev/pandas/issues/23023
             # block.shape is incorrect for "2D" ExtensionArrays
             # We can't, and don't need to, reshape.
 
-            # error: "ExtensionArray" has no attribute "reshape"
-            values = values.reshape(tuple((1,) + shape))  # type: ignore[attr-defined]
+            values = np.asarray(values).reshape(1, -1)
     return values
-
-
-def safe_reshape(arr: ArrayLike, new_shape: Shape) -> ArrayLike:
-    """
-    Reshape `arr` to have shape `new_shape`, unless it is an ExtensionArray,
-    in which case it will be returned unchanged (see gh-13012).
-
-    Parameters
-    ----------
-    arr : np.ndarray or ExtensionArray
-    new_shape : Tuple[int]
-
-    Returns
-    -------
-    np.ndarray or ExtensionArray
-    """
-    if not is_extension_array_dtype(arr.dtype):
-        # Note: this will include TimedeltaArray and tz-naive DatetimeArray
-        # TODO(EA2D): special case will be unnecessary with 2D EAs
-        arr = np.asarray(arr).reshape(new_shape)
-    return arr

pandas/core/internals/managers.py

@@ -69,10 +69,10 @@
     DatetimeTZBlock,
     ExtensionBlock,
     ObjectValuesExtensionBlock,
+    ensure_block_shape,
     extend_blocks,
     get_block_type,
     make_block,
-    safe_reshape,
 )
 from pandas.core.internals.ops import (
     blockwise_all,
@@ -1042,7 +1042,7 @@ def value_getitem(placement):
                 value = value.T
 
             if value.ndim == self.ndim - 1:
-                value = safe_reshape(value, (1,) + value.shape)
+                value = ensure_block_shape(value, ndim=2)
 
                 def value_getitem(placement):
                     return value
@@ -1167,7 +1167,7 @@ def insert(self, loc: int, item: Hashable, value, allow_duplicates: bool = False
             value = value.T
         elif value.ndim == self.ndim - 1 and not is_extension_array_dtype(value.dtype):
             # TODO(EA2D): special case not needed with 2D EAs
-            value = safe_reshape(value, (1,) + value.shape)
+            value = ensure_block_shape(value, ndim=2)
 
         block = make_block(values=value, ndim=self.ndim, placement=slice(loc, loc + 1))