Argument 1 to "apply" of "Series" has incompatible type "Callable[[Any], Union[int, NAType]]" (#724)

Update series.pyi and add test case

Author: ABizzinotto

pandas-stubs/core/series.pyi

@@ -851,7 +851,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     @overload
     def apply(
         self,
-        func: Callable[..., Scalar | Sequence | set | Mapping | None],
+        func: Callable[..., Scalar | Sequence | set | Mapping | NAType | None],
         convertDType: _bool = ...,
         args: tuple = ...,
         **kwds,

tests/test_series.py

@@ -35,6 +35,7 @@
 )
 import xarray as xr
 
+from pandas._libs.missing import NAType
 from pandas._libs.tslibs.timedeltas import Timedelta
 from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._typing import (
@@ -470,6 +471,8 @@ def get_depth(url: str) -> int:
     ss = s.astype(str)
     check(assert_type(ss.apply(get_depth), pd.Series), pd.Series, np.int64)
 
+    check(assert_type(s.apply(lambda x: pd.NA), pd.Series), pd.Series, NAType)
+
 
 def test_types_element_wise_arithmetic() -> None:
     s = pd.Series([0, 1, -10])