REF: call _block_shape from EABlock.make_block (#33308)

Author: jbrockmendel

pandas/core/internals/__init__.py

@@ -10,7 +10,6 @@
     IntBlock,
     ObjectBlock,
     TimeDeltaBlock,
-    _block_shape,
     _safe_reshape,
     make_block,
 )
@@ -31,7 +30,6 @@
     "TimeDeltaBlock",
     "_safe_reshape",
     "make_block",
-    "_block_shape",
     "BlockManager",
     "SingleBlockManager",
     "concatenate_block_managers",

pandas/core/internals/blocks.py

@@ -243,6 +243,8 @@ 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)
 
         return make_block(values, placement=placement, ndim=self.ndim)
 
@@ -354,13 +356,12 @@ def _split_op_result(self, result) -> List["Block"]:
             nbs = []
             for i, loc in enumerate(self.mgr_locs):
                 vals = result[i]
-                nv = _block_shape(vals, ndim=self.ndim)
-                block = self.make_block(values=nv, placement=[loc])
+                block = self.make_block(values=vals, placement=[loc])
                 nbs.append(block)
             return nbs
 
         if not isinstance(result, Block):
-            result = self.make_block(values=_block_shape(result, ndim=self.ndim))
+            result = self.make_block(result)
 
         return [result]
 
@@ -1264,9 +1265,6 @@ def take_nd(self, indexer, axis: int, new_mgr_locs=None, fill_value=lib.no_defau
     def diff(self, n: int, axis: int = 1) -> List["Block"]:
         """ return block for the diff of the values """
         new_values = algos.diff(self.values, n, axis=axis, stacklevel=7)
-        # We use block_shape for ExtensionBlock subclasses, which may call here
-        # via a super.
-        new_values = _block_shape(new_values, ndim=self.ndim)
         return [self.make_block(values=new_values)]
 
     def shift(self, periods: int, axis: int = 0, fill_value=None):
@@ -2254,7 +2252,7 @@ def concat_same_type(self, to_concat):
             values = values.astype(object, copy=False)
             placement = self.mgr_locs if self.ndim == 2 else slice(len(values))
 
-            return self.make_block(_block_shape(values, self.ndim), placement=placement)
+            return self.make_block(values, placement=placement)
         return super().concat_same_type(to_concat)
 
     def fillna(self, value, limit=None, inplace=False, downcast=None):
@@ -2420,7 +2418,6 @@ def f(mask, val, idx):
                 # TODO: allow EA once reshape is supported
                 values = values.reshape(shape)
 
-            values = _block_shape(values, ndim=self.ndim)
             return values
 
         if self.ndim == 2:
@@ -2660,9 +2657,7 @@ def concat_same_type(self, to_concat):
         )
         placement = self.mgr_locs if self.ndim == 2 else slice(len(values))
         # not using self.make_block_same_class as values can be object dtype
-        return self.make_block(
-            _block_shape(values, ndim=self.ndim), placement=placement
-        )
+        return self.make_block(values, placement=placement)
 
     def replace(
         self,
@@ -2771,16 +2766,15 @@ def _extend_blocks(result, blocks=None):
     return blocks
 
 
-def _block_shape(values, ndim=1, shape=None):
+def _block_shape(values: ArrayLike, ndim: int = 1) -> ArrayLike:
     """ guarantee the shape of the values to be at least 1 d """
     if values.ndim < ndim:
-        if shape is None:
-            shape = values.shape
-        if not is_extension_array_dtype(values):
-            # TODO: https://github.com/pandas-dev/pandas/issues/23023
+        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.
-            values = values.reshape(tuple((1,) + shape))
+            values = values.reshape(tuple((1,) + shape))  # type: ignore
     return values