From 1146cd678fc875c97fc821507ef974f7d53747ce Mon Sep 17 00:00:00 2001
From: Simon Hawkins <simonjayhawkins@gmail.com>
Date: Mon, 13 Apr 2020 12:04:01 +0100
Subject: [PATCH] TYP: use overload to refine return type of reset_index

---
 pandas/core/frame.py | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/pandas/core/frame.py b/pandas/core/frame.py
index d19f1a263f71a..23bea58de1a8e 100644
--- a/pandas/core/frame.py
+++ b/pandas/core/frame.py
@@ -32,6 +32,7 @@
     Type,
     Union,
     cast,
+    overload,
 )
 import warnings
 
@@ -4302,7 +4303,29 @@ def set_index(
         if not inplace:
             return frame
 
+    @overload
     def reset_index(
+        self,
+        level: Optional[Union[Hashable, Sequence[Hashable]]] = None,
+        drop: bool = False,
+        *,
+        col_level: Hashable = 0,
+        col_fill: Label = "",
+    ) -> "DataFrame":
+        ...
+
+    @overload
+    def reset_index(  # noqa: F811
+        self,
+        level: Optional[Union[Hashable, Sequence[Hashable]]] = None,
+        drop: bool = False,
+        inplace: bool = False,
+        col_level: Hashable = 0,
+        col_fill: Label = "",
+    ) -> Optional["DataFrame"]:
+        ...
+
+    def reset_index(  # noqa: F811
         self,
         level: Optional[Union[Hashable, Sequence[Hashable]]] = None,
         drop: bool = False,
@@ -6559,8 +6582,6 @@ def explode(self, column: Union[str, Tuple]) -> "DataFrame":
             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)
         result.index = self.index.take(result.index)