Fix replace, fillna allowing pd.NA. Allow list of arguments for .loc (#312)

* Fix replace, fillna allowing pd.NA.  Allow list of arguments for .loc

* add check and assert_type

Author: Dr-Irv

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=...,

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=...,

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)

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: