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..0aaab2081 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"]) + check(assert_type(frame.replace("N/A", pd.NA), pd.DataFrame), pd.DataFrame) 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: