diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 23c61773daf5a..260ff67742e97 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -4541,6 +4541,26 @@ def align( broadcast_axis=broadcast_axis, ) + @overload + def set_axis( + self, labels, axis: Axis = ..., inplace: Literal[False] = ... + ) -> DataFrame: + ... + + @overload + def set_axis(self, labels, axis: Axis, inplace: Literal[True]) -> None: + ... + + @overload + def set_axis(self, labels, *, inplace: Literal[True]) -> None: + ... + + @overload + def set_axis( + self, labels, axis: Axis = ..., inplace: bool = ... + ) -> Optional[DataFrame]: + ... + @Appender( """ Examples diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 050550f8add50..0b89cd83e4f76 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -24,6 +24,7 @@ Type, Union, cast, + overload, ) import warnings import weakref @@ -162,6 +163,8 @@ from pandas.io.formats.printing import pprint_thing if TYPE_CHECKING: + from typing import Literal + from pandas._libs.tslibs import BaseOffset from pandas.core.frame import DataFrame @@ -682,6 +685,28 @@ def _obj_with_exclusions(self: FrameOrSeries) -> FrameOrSeries: """ internal compat with SelectionMixin """ return self + @overload + def set_axis( + self: FrameOrSeries, labels, axis: Axis = ..., inplace: Literal[False] = ... + ) -> FrameOrSeries: + ... + + @overload + def set_axis( + self: FrameOrSeries, labels, axis: Axis, inplace: Literal[True] + ) -> None: + ... + + @overload + def set_axis(self: FrameOrSeries, labels, *, inplace: Literal[True]) -> None: + ... + + @overload + def set_axis( + self: FrameOrSeries, labels, axis: Axis = ..., inplace: bool = ... + ) -> Optional[FrameOrSeries]: + ... + def set_axis(self, labels, axis: Axis = 0, inplace: bool = False): """ Assign desired index to given axis. diff --git a/pandas/core/series.py b/pandas/core/series.py index 1f5a42265ec98..662c7abb33e33 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -20,6 +20,7 @@ Type, Union, cast, + overload, ) import warnings @@ -142,6 +143,8 @@ import pandas.plotting if TYPE_CHECKING: + from typing import Literal + from pandas._typing import ( TimedeltaConvertibleTypes, TimestampConvertibleTypes, @@ -4342,6 +4345,26 @@ def rename( else: return self._set_name(index, inplace=inplace) + @overload + def set_axis( + self, labels, axis: Axis = ..., inplace: Literal[False] = ... + ) -> Series: + ... + + @overload + def set_axis(self, labels, axis: Axis, inplace: Literal[True]) -> None: + ... + + @overload + def set_axis(self, labels, *, inplace: Literal[True]) -> None: + ... + + @overload + def set_axis( + self, labels, axis: Axis = ..., inplace: bool = ... + ) -> Optional[Series]: + ... + @Appender( """ Examples