CLN: use ensure_block_shape (pandas-dev#40792)

Author: jbrockmendel

pandas/core/internals/array_manager.py

@@ -479,8 +479,7 @@ def apply_with_block(self: T, f, align_keys=None, swap_axis=True, **kwargs) -> T
                 arr = arr._data  # type: ignore[attr-defined]
 
             if self.ndim == 2:
-                if isinstance(arr, np.ndarray):
-                    arr = np.atleast_2d(arr)
+                arr = ensure_block_shape(arr, 2)
                 block = new_block(arr, placement=slice(0, 1, 1), ndim=2)
             else:
                 block = new_block(arr, placement=slice(0, len(self), 1), ndim=1)
@@ -490,6 +489,7 @@ def apply_with_block(self: T, f, align_keys=None, swap_axis=True, **kwargs) -> T
                 applied = applied[0]
             arr = applied.values
             if self.ndim == 2 and arr.ndim == 2:
+                # 2D for np.ndarray or DatetimeArray/TimedeltaArray
                 assert len(arr) == 1
                 # error: Invalid index type "Tuple[int, slice]" for
                 # "Union[ndarray, ExtensionArray]"; expected type

pandas/core/internals/blocks.py

@@ -925,15 +925,6 @@ def setitem(self, indexer, value):
             # current dtype cannot store value, coerce to common dtype
             return self.coerce_to_target_dtype(value).setitem(indexer, value)
 
-        if self.dtype.kind in ["m", "M"]:
-            arr = self.values
-            if self.ndim > 1:
-                # Dont transpose with ndim=1 bc we would fail to invalidate
-                #  arr.freq
-                arr = arr.T
-            arr[indexer] = value
-            return self
-
         # value must be storable at this moment
         if is_extension_array_dtype(getattr(value, "dtype", None)):
             # We need to be careful not to allow through strings that
@@ -1697,6 +1688,19 @@ def iget(self, key):
         # TODO(EA2D): this can be removed if we ever have 2D EA
         return self.values.reshape(self.shape)[key]
 
+    def setitem(self, indexer, value):
+        if not self._can_hold_element(value):
+            # TODO: general case needs casting logic.
+            return self.astype(object).setitem(indexer, value)
+
+        values = self.values
+        if self.ndim > 1:
+            # Dont transpose with ndim=1 bc we would fail to invalidate
+            #  arr.freq
+            values = values.T
+        values[indexer] = value
+        return self
+
     def putmask(self, mask, new) -> list[Block]:
         mask = extract_bool_array(mask)
 

pandas/core/internals/managers.py

@@ -30,7 +30,7 @@
 
 from pandas.core.dtypes.cast import infer_dtype_from_scalar
 from pandas.core.dtypes.common import (
-    ensure_int64,
+    ensure_platform_int,
     is_dtype_equal,
     is_extension_array_dtype,
     is_list_like,
@@ -346,6 +346,13 @@ def unpickle_block(values, mgr_locs, ndim: int):
             state = state[3]["0.14.1"]
             self.axes = [ensure_index(ax) for ax in state["axes"]]
             ndim = len(self.axes)
+
+            for blk in state["blocks"]:
+                vals = blk["values"]
+                # older versions may hold e.g. DatetimeIndex instead of DTA
+                vals = extract_array(vals, extract_numpy=True)
+                blk["values"] = ensure_block_shape(vals, ndim=ndim)
+
             self.blocks = tuple(
                 unpickle_block(b["values"], b["mgr_locs"], ndim=ndim)
                 for b in state["blocks"]
@@ -1975,8 +1982,7 @@ def _preprocess_slice_or_indexer(
             dtype = getattr(slice_or_indexer, "dtype", None)
             raise TypeError(type(slice_or_indexer), dtype)
 
-        # TODO: np.intp?
-        indexer = ensure_int64(slice_or_indexer)
+        indexer = ensure_platform_int(slice_or_indexer)
         if not allow_fill:
             indexer = maybe_convert_indices(indexer, length)
         return "fancy", indexer, len(indexer)

pandas/tests/internals/test_internals.py

@@ -37,7 +37,10 @@
     SingleBlockManager,
     make_block,
 )
-from pandas.core.internals.blocks import new_block
+from pandas.core.internals.blocks import (
+    ensure_block_shape,
+    new_block,
+)
 
 # this file contains BlockManager specific tests
 # TODO(ArrayManager) factor out interleave_dtype tests
@@ -138,6 +141,7 @@ def create_block(typestr, placement, item_shape=None, num_offset=0, maker=new_bl
         tz = m.groups()[0]
         assert num_items == 1, "must have only 1 num items for a tz-aware"
         values = DatetimeIndex(np.arange(N) * 1e9, tz=tz)._data
+        values = ensure_block_shape(values, ndim=len(shape))
     elif typestr in ("timedelta", "td", "m8[ns]"):
         values = (mat * 1).astype("m8[ns]")
     elif typestr in ("category",):

pandas/tests/series/methods/test_isin.py

@@ -59,7 +59,7 @@ def test_isin_with_i8(self):
         tm.assert_series_equal(result, expected)
 
         # fails on dtype conversion in the first place
-        result = s.isin(s[0:2].values.astype("datetime64[D]"))
+        result = s.isin(np.asarray(s[0:2].values).astype("datetime64[D]"))
         tm.assert_series_equal(result, expected)
 
         result = s.isin([s[1]])