diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 539275c7ff617..3607fbea594c6 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -34,6 +34,7 @@ Type, Union, cast, + overload, ) import warnings @@ -155,6 +156,8 @@ import pandas.plotting if TYPE_CHECKING: + from typing import Literal + from pandas.core.groupby.generic import DataFrameGroupBy from pandas.io.formats.style import Styler @@ -4692,6 +4695,30 @@ def set_index( if not inplace: return frame + @overload + # https://github.com/python/mypy/issues/6580 + # Overloaded function signatures 1 and 2 overlap with incompatible return types + def reset_index( # type: ignore[misc] + self, + level: Optional[Union[Hashable, Sequence[Hashable]]] = ..., + drop: bool = ..., + inplace: Literal[False] = ..., + col_level: Hashable = ..., + col_fill: Label = ..., + ) -> DataFrame: + ... + + @overload + def reset_index( + self, + level: Optional[Union[Hashable, Sequence[Hashable]]] = ..., + drop: bool = ..., + inplace: Literal[True] = ..., + col_level: Hashable = ..., + col_fill: Label = ..., + ) -> None: + ... + def reset_index( self, level: Optional[Union[Hashable, Sequence[Hashable]]] = None, @@ -7156,8 +7183,6 @@ def explode( raise ValueError("columns must be unique") df = self.reset_index(drop=True) - # TODO: use overload to refine return type of reset_index - assert df is not None # needed for mypy result = df[column].explode() result = df.drop([column], axis=1).join(result) if ignore_index: