GH1214 Improve typehinting for DataFrame.unstack

Author: loicdiridollou

pandas-stubs/core/frame.pyi

@@ -1300,7 +1300,8 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     def unstack(
         self,
         level: Level = ...,
-        fill_value: int | _str | dict | None = ...,
+        fill_value: Scalar | None = ...,
+        sort: _bool = ...,
     ) -> Self | Series: ...
     def melt(
         self,

tests/test_frame.py

@@ -4339,3 +4339,49 @@ def test_df_loc_dict() -> None:
 
     df.iloc[0] = {"X": 0}
     check(assert_type(df, pd.DataFrame), pd.DataFrame)
+
+
+def test_unstack() -> None:
+    """Test different types of argument for `fill_value` in DataFrame.unstack."""
+    df = pd.DataFrame(
+        [
+            ["a", "b", pd.Timestamp(2021, 3, 2)],
+            ["a", "a", pd.Timestamp(2023, 4, 2)],
+            ["b", "b", pd.Timestamp(2024, 3, 2)],
+        ]
+    ).set_index([0, 1])
+
+    check(assert_type(df.unstack(0), pd.DataFrame | pd.Series), pd.DataFrame)
+    check(
+        assert_type(
+            df.unstack(1, fill_value=pd.Timestamp(2023, 4, 5)), pd.DataFrame | pd.Series
+        ),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.unstack(1, fill_value=0.0), pd.DataFrame | pd.Series),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.unstack(1, fill_value=1), pd.DataFrame | pd.Series), pd.DataFrame
+    )
+    check(
+        assert_type(df.unstack(1, fill_value="string"), pd.DataFrame | pd.Series),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(df.unstack(0, sort=False), pd.DataFrame | pd.Series), pd.DataFrame
+    )
+    check(
+        assert_type(
+            df.unstack(1, fill_value=pd.Timestamp(2023, 4, 5), sort=True),
+            pd.DataFrame | pd.Series,
+        ),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            df.unstack(1, fill_value=0.0, sort=False), pd.DataFrame | pd.Series
+        ),
+        pd.DataFrame,
+    )