diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 1f8a2648..10d6a232 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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, diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi index 456d0697..c8edf865 100644 --- a/pandas-stubs/core/generic.pyi +++ b/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: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 85943841..3e607e65 100644 --- a/pandas-stubs/core/series.pyi +++ b/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__( diff --git a/tests/test_frame.py b/tests/test_frame.py index 7d4890ee..a7536b91 100644 --- a/tests/test_frame.py +++ b/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]}) diff --git a/tests/test_series.py b/tests/test_series.py index 9edc8e06..ac6d897e 100644 --- a/tests/test_series.py +++ b/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