From a860d0f58ef2d2b06c4a06370b7a137c74ce7046 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= Date: Thu, 15 May 2025 20:29:40 -0400 Subject: [PATCH 1/2] GH1214 Improve typehinting for DataFrame.unstack --- pandas-stubs/core/frame.pyi | 3 ++- tests/test_frame.py | 46 +++++++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 1f8a2648..fb2c1c9b 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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, diff --git a/tests/test_frame.py b/tests/test_frame.py index 7d4890ee..3fd76bef 100644 --- a/tests/test_frame.py +++ b/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, + ) From 23adb71f1502d52a5d43b5b9b0fecc62334f0712 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= Date: Thu, 15 May 2025 20:46:15 -0400 Subject: [PATCH 2/2] GH1214 Fix test --- tests/test_frame.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index 3fd76bef..9b666bbc 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -4350,6 +4350,20 @@ def test_unstack() -> None: ["b", "b", pd.Timestamp(2024, 3, 2)], ] ).set_index([0, 1]) + df_sr = pd.DataFrame( + [ + ["a", "b", "abc"], + ["a", "a", "def"], + ["b", "b", "ghi"], + ] + ).set_index([0, 1]) + df_flt = pd.DataFrame( + [ + ["a", "b", 1], + ["a", "a", 12], + ["b", "b", 14], + ] + ).set_index([0, 1]) check(assert_type(df.unstack(0), pd.DataFrame | pd.Series), pd.DataFrame) check( @@ -4359,14 +4373,15 @@ def test_unstack() -> None: pd.DataFrame, ) check( - assert_type(df.unstack(1, fill_value=0.0), pd.DataFrame | pd.Series), + assert_type(df_flt.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 + assert_type(df_flt.unstack(1, fill_value=1), pd.DataFrame | pd.Series), + pd.DataFrame, ) check( - assert_type(df.unstack(1, fill_value="string"), pd.DataFrame | pd.Series), + assert_type(df_sr.unstack(1, fill_value="string"), pd.DataFrame | pd.Series), pd.DataFrame, ) check( @@ -4381,7 +4396,7 @@ def test_unstack() -> None: ) check( assert_type( - df.unstack(1, fill_value=0.0, sort=False), pd.DataFrame | pd.Series + df_flt.unstack(1, fill_value=0.0, sort=False), pd.DataFrame | pd.Series ), pd.DataFrame, )