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

[bersbersbers]

To Reproduce

1. Provide a minimal runnable pandas example that is not properly checked by the stubs.

import pandas as pd

# these values are allowed
allowed = ["A", "B", "C"]

df = pd.DataFrame()

# we add a few that are not allowed
df["Column"] = ["A", "B", "C", "D"]

# replace values by index, and NA if not allowed
df["Label"] = df.Column.apply(lambda x: allowed.index(x) if x in allowed else pd.NA)

print(df.Label)

2. Indicate which type checker you are using (mypy or pyright).
   mypy 1.2.0 (compiled: yes)

3. Show the error message received from that type checker while checking your example.

bug.py:12: error: Argument 1 to "apply" of "Series" has incompatible type "Callable[[Any], Union[int, NAType]]"; expected "Callable[..., Union[str, bytes, date, datetime, timedelta, datetime64, timedelta64, bool, int, float, Timestamp, Timedelta, integer[Any], floating[Any], complex, Sequence[Any], Set[Any], Mapping[Any, Any]]]"  [arg-type]
bug.py:12: error: Incompatible return value type (got "Union[int, NAType]", expected "Union[str, bytes, date, datetime, timedelta, datetime64, timedelta64, bool, int, float, Timestamp, Timedelta, integer[Any], floating[Any], complex, Sequence[Any], Set[Any], Mapping[Any, Any]]")  [return-value]
Found 2 errors in 1 file (checked 1 source file)

Please complete the following information:

• OS: Windows 11
• OS Version: 22H2
• python version: 3.10.11
• version of type checker: 1.2.0
• version of installed pandas-stubs: 2.0.1.230501

Related: #262?

[twoertwein]

Probably need to add NAType here

pandas-stubs/pandas-stubs/core/series.pyi

Line 784 in 6c0eed6

func: Callable[..., Scalar | Sequence | set | Mapping | None],

or even to Scalar / S1.

[Dr-Irv]

Probably need to add NAType here

Correct

pandas-stubs/pandas-stubs/core/series.pyi

Line 784 in 6c0eed6

func: Callable[..., Scalar | Sequence | set | Mapping | None],

or even to Scalar / S1.

I'm wary (as I wrote #262 (comment)) of doing the latter because it may make some of our existing declarations too wide.

So better to make the change as suggested above to handle this case.