REF: stricter typing in Manager.take (#51478)

Author: jbrockmendel

pandas/core/generic.py

@@ -3895,6 +3895,7 @@ class  max_speed
 
         return self._take(indices, axis)
 
+    @final
     def _take(
         self: NDFrameT,
         indices,
@@ -3915,6 +3916,11 @@ def _take(
                 and is_range_indexer(indices, len(self))
             ):
                 return self.copy(deep=None)
+        else:
+            # We can get here with a slice via DataFrame.__geittem__
+            indices = np.arange(
+                indices.start, indices.stop, indices.step, dtype=np.intp
+            )
 
         new_data = self._mgr.take(
             indices,

pandas/core/internals/array_manager.py

@@ -630,21 +630,18 @@ def _reindex_indexer(
 
     def take(
         self: T,
-        indexer,
+        indexer: npt.NDArray[np.intp],
         axis: AxisInt = 1,
         verify: bool = True,
         convert_indices: bool = True,
     ) -> T:
         """
         Take items along any axis.
         """
-        axis = self._normalize_axis(axis)
+        assert isinstance(indexer, np.ndarray), type(indexer)
+        assert indexer.dtype == np.intp, indexer.dtype
 
-        indexer = (
-            np.arange(indexer.start, indexer.stop, indexer.step, dtype="int64")
-            if isinstance(indexer, slice)
-            else np.asanyarray(indexer, dtype="int64")
-        )
+        axis = self._normalize_axis(axis)
 
         if not indexer.ndim == 1:
             raise ValueError("indexer should be 1-dimensional")

pandas/core/internals/managers.py

@@ -907,15 +907,15 @@ def _make_na_block(
 
     def take(
         self: T,
-        indexer,
+        indexer: npt.NDArray[np.intp],
         axis: AxisInt = 1,
         verify: bool = True,
         convert_indices: bool = True,
     ) -> T:
         """
         Take items along any axis.
 
-        indexer : np.ndarray or slice
+        indexer : np.ndarray[np.intp]
         axis : int, default 1
         verify : bool, default True
             Check that all entries are between 0 and len(self) - 1, inclusive.
@@ -927,12 +927,8 @@ def take(
         -------
         BlockManager
         """
-        # We have 6 tests that get here with a slice
-        indexer = (
-            np.arange(indexer.start, indexer.stop, indexer.step, dtype=np.intp)
-            if isinstance(indexer, slice)
-            else np.asanyarray(indexer, dtype=np.intp)
-        )
+        assert isinstance(indexer, np.ndarray), type(indexer)
+        assert indexer.dtype == np.intp, indexer.dtype
 
         n = self.shape[axis]
         if convert_indices:

pandas/tests/internals/test_internals.py

@@ -1003,13 +1003,15 @@ def assert_take_ok(mgr, axis, indexer):
 
         for ax in range(mgr.ndim):
             # take/fancy indexer
-            assert_take_ok(mgr, ax, indexer=[])
-            assert_take_ok(mgr, ax, indexer=[0, 0, 0])
-            assert_take_ok(mgr, ax, indexer=list(range(mgr.shape[ax])))
+            assert_take_ok(mgr, ax, indexer=np.array([], dtype=np.intp))
+            assert_take_ok(mgr, ax, indexer=np.array([0, 0, 0], dtype=np.intp))
+            assert_take_ok(
+                mgr, ax, indexer=np.array(list(range(mgr.shape[ax])), dtype=np.intp)
+            )
 
             if mgr.shape[ax] >= 3:
-                assert_take_ok(mgr, ax, indexer=[0, 1, 2])
-                assert_take_ok(mgr, ax, indexer=[-1, -2, -3])
+                assert_take_ok(mgr, ax, indexer=np.array([0, 1, 2], dtype=np.intp))
+                assert_take_ok(mgr, ax, indexer=np.array([-1, -2, -3], dtype=np.intp))
 
     @pytest.mark.parametrize("mgr", MANAGERS)
     @pytest.mark.parametrize("fill_value", [None, np.nan, 100.0])