From e47a276e7adc2af237a5783ae551c809cf28313f Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Sat, 17 Sep 2022 16:05:25 -0400 Subject: [PATCH 1/2] Fix replace, fillna allowing pd.NA. Allow list of arguments for .loc --- pandas-stubs/core/frame.pyi | 16 +++++++++------- pandas-stubs/core/series.pyi | 13 +++++++------ tests/test_frame.py | 16 ++++++++++++++++ tests/test_series.py | 2 ++ 4 files changed, 34 insertions(+), 13 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index ee625bbef..9e70fbeff 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -40,6 +40,7 @@ from pandas.core.window.rolling import ( ) import xarray as xr +from pandas._libs.missing import NAType from pandas._typing import ( S1, AggFuncType, @@ -157,7 +158,8 @@ class _LocIndexerFrame(_LocIndexer): self, idx: MaskType | StrLike - | tuple[MaskType | Index | Sequence[ScalarT] | Scalar | slice, ...], + | tuple[MaskType | Index | Sequence[ScalarT] | Scalar | slice, ...] + | list[ScalarT], value: S1 | ArrayLike | Series | DataFrame, ) -> None: ... @overload @@ -617,7 +619,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def fillna( self, - value: Scalar | dict | Series | DataFrame | None = ..., + value: Scalar | NAType | dict | Series | DataFrame | None = ..., method: FillnaOptions | None = ..., axis: AxisType | None = ..., limit: int = ..., @@ -628,7 +630,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def fillna( self, - value: Scalar | dict | Series | DataFrame | None = ..., + value: Scalar | NAType | dict | Series | DataFrame | None = ..., method: FillnaOptions | None = ..., axis: AxisType | None = ..., limit: int = ..., @@ -639,7 +641,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def fillna( self, - value: Scalar | dict | Series | DataFrame | None = ..., + value: Scalar | NAType | dict | Series | DataFrame | None = ..., method: FillnaOptions | None = ..., axis: AxisType | None = ..., inplace: _bool | None = ..., @@ -650,7 +652,7 @@ class DataFrame(NDFrame, OpsMixin): def replace( self, to_replace=..., - value: Scalar | Sequence | Mapping | Pattern | None = ..., + value: Scalar | NAType | Sequence | Mapping | Pattern | None = ..., limit: int | None = ..., regex=..., method: ReplaceMethod = ..., @@ -661,7 +663,7 @@ class DataFrame(NDFrame, OpsMixin): def replace( self, to_replace=..., - value: Scalar | Sequence | Mapping | Pattern | None = ..., + value: Scalar | NAType | Sequence | Mapping | Pattern | None = ..., inplace: Literal[False] = ..., limit: int | None = ..., regex=..., @@ -671,7 +673,7 @@ class DataFrame(NDFrame, OpsMixin): def replace( self, to_replace=..., - value: Scalar | Sequence | Mapping | Pattern | None = ..., + value: Scalar | NAType | Sequence | Mapping | Pattern | None = ..., inplace: _bool | None = ..., limit: int | None = ..., regex=..., diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index e64e56fc4..5e12e20cf 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -56,6 +56,7 @@ from pandas.core.window.rolling import ( ) import xarray as xr +from pandas._libs.missing import NAType from pandas._typing import ( S1, AggFuncTypeBase, @@ -772,7 +773,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def fillna( self, - value: Scalar | dict | Series[S1] | DataFrame | None = ..., + value: Scalar | NAType | dict | Series[S1] | DataFrame | None = ..., method: FillnaOptions | None = ..., axis: SeriesAxisType = ..., limit: int | None = ..., @@ -783,7 +784,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def fillna( self, - value: Scalar | dict | Series[S1] | DataFrame | None = ..., + value: Scalar | NAType | dict | Series[S1] | DataFrame | None = ..., method: FillnaOptions | None = ..., axis: SeriesAxisType = ..., *, @@ -793,7 +794,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def fillna( self, - value: Scalar | dict | Series[S1] | DataFrame | None = ..., + value: Scalar | NAType | dict | Series[S1] | DataFrame | None = ..., method: FillnaOptions | None = ..., axis: SeriesAxisType = ..., inplace: _bool = ..., @@ -804,7 +805,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def replace( self, to_replace: _str | list | dict | Series[S1] | float | None = ..., - value: Scalar | dict | list | _str | None = ..., + value: Scalar | NAType | dict | list | _str | None = ..., inplace: Literal[False] = ..., limit: int | None = ..., regex=..., @@ -814,7 +815,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def replace( self, to_replace: _str | list | dict | Series[S1] | float | None = ..., - value: Scalar | dict | list | _str | None = ..., + value: Scalar | NAType | dict | list | _str | None = ..., limit: int | None = ..., regex=..., method: ReplaceMethod = ..., @@ -825,7 +826,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): def replace( self, to_replace: _str | list | dict | Series[S1] | float | None = ..., - value: Scalar | dict | list | _str | None = ..., + value: Scalar | NAType | dict | list | _str | None = ..., inplace: _bool = ..., limit: int | None = ..., regex=..., diff --git a/tests/test_frame.py b/tests/test_frame.py index c1e66c3fe..ff3fb1922 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1781,3 +1781,19 @@ def cond2(x: pd.DataFrame) -> pd.DataFrame: cond3 = pd.DataFrame({"a": [True, True, False], "b": [False, False, False]}) check(assert_type(df.where(cond3), pd.DataFrame), pd.DataFrame) + + +def test_setitem_loc() -> None: + # GH 254 + df = pd.DataFrame.from_dict( + {view: (True, True, True) for view in ["A", "B", "C"]}, orient="index" + ) + df.loc[["A", "C"]] = False + my_arr = ["A", "C"] + df.loc[my_arr] = False + + +def test_replace_na() -> None: + # GH 262 + frame = pd.DataFrame(["N/A", "foo", "bar"]) + frame = frame.replace("N/A", pd.NA) diff --git a/tests/test_series.py b/tests/test_series.py index 175dd638d..25696bf4c 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -209,6 +209,8 @@ def test_types_fillna() -> None: assert assert_type(s.fillna(method="bfill", inplace=True), None) is None check(assert_type(s.fillna(method="pad"), pd.Series), pd.Series) check(assert_type(s.fillna(method="ffill", limit=1), pd.Series), pd.Series) + # GH 263 + check(assert_type(s.fillna(pd.NA), pd.Series), pd.Series) def test_types_sort_index() -> None: From bae2c58b9111414f72cb148e46586568f4350772 Mon Sep 17 00:00:00 2001 From: Irv Lustig Date: Sat, 17 Sep 2022 16:08:04 -0400 Subject: [PATCH 2/2] add check and assert_type --- tests/test_frame.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_frame.py b/tests/test_frame.py index ff3fb1922..0aaab2081 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1796,4 +1796,4 @@ def test_setitem_loc() -> None: def test_replace_na() -> None: # GH 262 frame = pd.DataFrame(["N/A", "foo", "bar"]) - frame = frame.replace("N/A", pd.NA) + check(assert_type(frame.replace("N/A", pd.NA), pd.DataFrame), pd.DataFrame)