take method on NDDataFrame (#1209)

* take method on NDDataFrame

* add axis test cases & refactor

---------

Co-authored-by: Ritsuki Yamada <[email protected]>

Author: ritsuki1227

pandas-stubs/core/frame.pyi

@@ -2298,7 +2298,6 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     ) -> Series: ...
     def swapaxes(self, axis1: Axis, axis2: Axis, copy: _bool = ...) -> Self: ...
     def tail(self, n: int = ...) -> Self: ...
-    def take(self, indices: list, axis: Axis = ..., **kwargs: Any) -> Self: ...
     @overload
     def to_json(
         self,

pandas-stubs/core/generic.pyi

@@ -46,6 +46,7 @@ from pandas._typing import (
     P,
     StorageOptions,
     T,
+    TakeIndexer,
     TimedeltaConvertibleTypes,
     TimeGrouperOrigin,
     TimestampConvention,
@@ -418,3 +419,5 @@ class NDFrame(indexing.IndexingMixin):
         offset: TimedeltaConvertibleTypes | None = ...,
         group_keys: _bool = ...,
     ) -> DatetimeIndexResampler[Self]: ...
+    @final
+    def take(self, indices: TakeIndexer, axis: Axis = ..., **kwargs: Any) -> Self: ...

pandas-stubs/core/series.pyi

@@ -450,9 +450,6 @@ class Series(IndexOpsMixin[S1], NDFrame):
     def __array__(self, dtype=...) -> np.ndarray: ...
     @property
     def axes(self) -> list: ...
-    def take(
-        self, indices: Sequence, axis: AxisIndex = ..., **kwargs: Any
-    ) -> Series[S1]: ...
     def __getattr__(self, name: _str) -> S1: ...
     @overload
     def __getitem__(

tests/test_frame.py

@@ -2984,6 +2984,17 @@ def test_iloc_tuple() -> None:
     df = df.iloc[0:2,]
 
 
+def test_take() -> None:
+    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    check(assert_type(df.take([0, 1]), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([np.int64(0), np.int64(1)]), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take(np.array([0, 1])), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([0, 1], "index"), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([0, 1], 0), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([0, 1], "columns"), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.take([0, 1], 1), pd.DataFrame), pd.DataFrame)
+
+
 def test_set_columns() -> None:
     # GH 73
     df = pd.DataFrame({"a": [1, 2, 3], "b": [0.0, 1, 1]})

tests/test_series.py

@@ -1715,6 +1715,16 @@ def test_iloc_getitem_ndarray() -> None:
     check(assert_type(values_s.iloc[indices_u64], pd.Series), pd.Series)
 
 
+def test_take() -> None:
+    s = pd.Series(np.arange(10), name="a")
+    check(assert_type(s.take([0, 1]), pd.Series), pd.Series)
+    check(
+        assert_type(s.take([np.int64(0), np.int64(1)]), pd.Series),
+        pd.Series,
+    )
+    check(assert_type(s.take(np.array([0, 1])), pd.Series), pd.Series)
+
+
 def test_iloc_setitem_ndarray() -> None:
     # GH 85
     # GH 86