REF: Implement core._algos (pandas-dev#32767)

Author: jbrockmendel

pandas/core/array_algos/__init__.py

@@ -0,0 +1,9 @@
+"""
+core.array_algos is for algorithms that operate on ndarray and ExtensionArray.
+These should:
+
+- Assume that any Index, Series, or DataFrame objects have already been unwrapped.
+- Assume that any list arguments have already been cast to ndarray/EA.
+- Not depend on Index, Series, or DataFrame, nor import any of these.
+- May dispatch to ExtensionArray methods, but should not import from core.arrays.
+"""

pandas/core/array_algos/transforms.py

@@ -0,0 +1,33 @@
+"""
+transforms.py is for shape-preserving functions.
+"""
+
+import numpy as np
+
+from pandas.core.dtypes.common import ensure_platform_int
+
+
+def shift(values: np.ndarray, periods: int, axis: int, fill_value) -> np.ndarray:
+    new_values = values
+
+    # make sure array sent to np.roll is c_contiguous
+    f_ordered = values.flags.f_contiguous
+    if f_ordered:
+        new_values = new_values.T
+        axis = new_values.ndim - axis - 1
+
+    if np.prod(new_values.shape):
+        new_values = np.roll(new_values, ensure_platform_int(periods), axis=axis)
+
+    axis_indexer = [slice(None)] * values.ndim
+    if periods > 0:
+        axis_indexer[axis] = slice(None, periods)
+    else:
+        axis_indexer[axis] = slice(periods, None)
+    new_values[tuple(axis_indexer)] = fill_value
+
+    # restore original order
+    if f_ordered:
+        new_values = new_values.T
+
+    return new_values

pandas/core/arrays/datetimelike.py

@@ -40,6 +40,7 @@
 
 from pandas.core import missing, nanops, ops
 from pandas.core.algorithms import checked_add_with_arr, take, unique1d, value_counts
+from pandas.core.array_algos.transforms import shift
 from pandas.core.arrays.base import ExtensionArray, ExtensionOpsMixin
 import pandas.core.common as com
 from pandas.core.construction import array, extract_array
@@ -773,26 +774,7 @@ def shift(self, periods=1, fill_value=None, axis=0):
 
         fill_value = self._unbox_scalar(fill_value)
 
-        new_values = self._data
-
-        # make sure array sent to np.roll is c_contiguous
-        f_ordered = new_values.flags.f_contiguous
-        if f_ordered:
-            new_values = new_values.T
-            axis = new_values.ndim - axis - 1
-
-        new_values = np.roll(new_values, periods, axis=axis)
-
-        axis_indexer = [slice(None)] * self.ndim
-        if periods > 0:
-            axis_indexer[axis] = slice(None, periods)
-        else:
-            axis_indexer[axis] = slice(periods, None)
-        new_values[tuple(axis_indexer)] = fill_value
-
-        # restore original order
-        if f_ordered:
-            new_values = new_values.T
+        new_values = shift(self._data, periods, axis, fill_value)
 
         return type(self)._simple_new(new_values, dtype=self.dtype)
 

pandas/core/internals/blocks.py

@@ -29,7 +29,6 @@
 from pandas.core.dtypes.common import (
     _NS_DTYPE,
     _TD_DTYPE,
-    ensure_platform_int,
     is_bool_dtype,
     is_categorical,
     is_categorical_dtype,
@@ -66,6 +65,7 @@
 )
 
 import pandas.core.algorithms as algos
+from pandas.core.array_algos.transforms import shift
 from pandas.core.arrays import (
     Categorical,
     DatetimeArray,
@@ -1316,25 +1316,7 @@ def shift(self, periods, axis: int = 0, fill_value=None):
         # that, handle boolean etc also
         new_values, fill_value = maybe_upcast(self.values, fill_value)
 
-        # make sure array sent to np.roll is c_contiguous
-        f_ordered = new_values.flags.f_contiguous
-        if f_ordered:
-            new_values = new_values.T
-            axis = new_values.ndim - axis - 1
-
-        if np.prod(new_values.shape):
-            new_values = np.roll(new_values, ensure_platform_int(periods), axis=axis)
-
-        axis_indexer = [slice(None)] * self.ndim
-        if periods > 0:
-            axis_indexer[axis] = slice(None, periods)
-        else:
-            axis_indexer[axis] = slice(periods, None)
-        new_values[tuple(axis_indexer)] = fill_value
-
-        # restore original order
-        if f_ordered:
-            new_values = new_values.T
+        new_values = shift(new_values, periods, axis, fill_value)
 
         return [self.make_block(new_values)]