Skip to content

Fix series.map overloads #960

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions pandas-stubs/_typing.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -568,6 +568,8 @@ S2 = TypeVar(
| BaseOffset,
)

SN = TypeVar("SN", bound=bool | int | float | complex)

IndexingInt: TypeAlias = (
int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8
)
Expand Down
9 changes: 8 additions & 1 deletion pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ from pandas._libs.tslibs.nattype import NaTType
from pandas._typing import (
S1,
S2,
SN,
AggFuncTypeBase,
AggFuncTypeDictFrame,
AggFuncTypeSeriesToFrame,
Expand Down Expand Up @@ -916,10 +917,16 @@ class Series(IndexOpsMixin[S1], NDFrame):
@overload
def map(
self,
arg: Callable[[S1], S2 | NAType] | Mapping[S1, S2] | Series[S2],
arg: Callable[[SN], S2 | NAType] | Mapping[SN, S2] | Series[S2],
na_action: Literal["ignore"] = ...,
) -> Series[S2]: ...
@overload
def map(
self,
arg: Callable[[S1], S2 | NAType] | Mapping[S1, S2] | Series[S2],
na_action: Literal["ignore"],
) -> Series[S2]: ...
@overload
def map(
self,
arg: Callable[[S1 | NAType], S2 | NAType] | Mapping[S1, S2] | Series[S2],
Expand Down
37 changes: 37 additions & 0 deletions tests/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -3276,6 +3276,7 @@ def test_map_na() -> None:

mapping = {1: "a", 2: "b", 3: "c"}
check(assert_type(s.map(mapping, na_action=None), "pd.Series[str]"), pd.Series, str)
check(assert_type(s.map(mapping), "pd.Series[str]"), pd.Series, str)

def callable(x: int | NAType) -> str | NAType:
if isinstance(x, int):
Expand All @@ -3285,9 +3286,45 @@ def callable(x: int | NAType) -> str | NAType:
check(
assert_type(s.map(callable, na_action=None), "pd.Series[str]"), pd.Series, str
)
check(assert_type(s.map(callable), "pd.Series[str]"), pd.Series, str)

series = pd.Series(["a", "b", "c"])
check(assert_type(s.map(series, na_action=None), "pd.Series[str]"), pd.Series, str)
check(assert_type(s.map(series), "pd.Series[str]"), pd.Series, str)

s2: pd.Series[float] = pd.Series([1.0, pd.NA, 3.0])
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We generally don't use the annotations in the tests. So you can do

s2 = pd.Series([1.0, pd.NA, 3.0], dtype=float)
# OR
s2 = pd.Series([1.0, pd.NA, 3.0]).astype(pd.Float64Dtype())

Then the type checkers should recognize s2 as Series[float]

Note that internally, pd.Series([1.0, pd.NA, 3.0]) has dtype object

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This first actually gives me an error TypeError: float() argument must be a string or a real number, not 'NAType'. While the second version does work, i cant get it to work for the integer case of s = pd.Series([1, pd.NA, 3]). While that does work at the construction point and for inferring s: Series[int], the values that are passed into callable at s.map(callable) are actually floats.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For Series[int], you can use s = pd.Series([1, pd.NA, 2], dtype="Int64")

Sorry about my error on the dtype=float.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

# Step 1: Create the Series with pd.NA
s = pd.Series([1, pd.NA, 3], dtype="Int64")

print(s)
print(s.dtype)

# Define the callable function
def callable(x: int | NAType) -> str | NAType:
    print(x, type(x))
    if pd.isna(x):
        return pd.NA
    return str(x)

print(s)
print(s.map(callable, na_action=None))

actually still gives me this:

$ python testing.py
0       1
1    <NA>
2       3
dtype: Int64
Int64
0       1
1    <NA>
2       3
dtype: Int64
1.0 <class 'float'>
nan <class 'float'>
3.0 <class 'float'>
0     1.0
1    <NA>
2     3.0
dtype: object

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a known pandas issue. See pandas-dev/pandas#57189


def callable2(x: float) -> float:
Copy link
Contributor Author

@JanEricNitschke JanEricNitschke Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Float (and int, complex, bool) callables pass na_action="ignore" and default na_action despite no explicit None handling, but fail the explicit na_action=None case as it is assumed that this is explicitly specified for rigorous type checking.

return x + 1

check(
assert_type(s2.map(callable2, na_action="ignore"), "pd.Series[float]"),
pd.Series,
float,
)
check(
assert_type(s2.map(callable2), "pd.Series[float]"),
pd.Series,
float,
)
if TYPE_CHECKING_INVALID_USAGE:
s2.map(callable2, na_action=None) # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType]

s3: pd.Series[str] = pd.Series(["A", pd.NA, "C"])
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similar comment. Use s3 = pd.Series(["A", pd.NA, "C"], dtype=str)


def callable3(x: str) -> str:
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-numerical functions pass only the explicit na_action="ignore" but not the na_action=None case (explicit or implicit) because these are likely to cause runtime errors on nan-values.

return x.lower()

check(
assert_type(s3.map(callable3, na_action="ignore"), "pd.Series[str]"),
pd.Series,
str,
)
if TYPE_CHECKING_INVALID_USAGE:
s3.map(
callable3, na_action=None # type: ignore[arg-type] # pyright: ignore[reportCallIssue, reportArgumentType]
)
s3.map(callable3) # type: ignore[type-var] # pyright: ignore[reportCallIssue, reportArgumentType]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a little concerned that this won't pass the type checker. A user should be able to do:

s = pd.Series({"a", "b", "c"], dtype=str)
s.map(callable3)

That will execute correctly and should pass the type checker.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unfortunately i dont think the type system supports something like this at the moment.

All that s.map(callable3) sees is that the entries of s are string or none values and that none-values will be passed into callable3 as they are. So it has to be able to handle them. The type system does not know if there actually are non-values in there or not, only that there might be. Most of the times the user probably wont know either if the data is loaded from an external file that comes from a sensor or a data repository.

If the user does know more than the type system he could either specify na_action="ignore", which he should do anyway if the function does no handle the case explicitly, or add a type ignore to the line as one normally does when you have more information than can be contained in the type system (or understood by the type checkers).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I understand your point of view. But the issue here is that we need to support the "normal" usage of pandas. The type system can't keep track of what is inside of a Series or DataFrame. For example, if you have a DataFrame df with columns "a", "b" then df["a"] is Series[Any] because we can't track the type inside the DataFrame.

The default of the na_action argument in Series.map() is None, which means that people who know that their series has no non-null elements would now have to handle the case in the declared typing of callable3() to handle nulls, but that will likely break existing code.

We've had other cases where we were too restrictive in the stubs, and people complained, so we had to pull back.

See https://github.com/pandas-dev/pandas-stubs/blob/main/docs/philosophy.md#narrow-vs-wide-arguments for a discussion of this issue. IMHO, you're proposing that the stubs will be too narrow.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, i think im just too much of a typing purist in this case. Because, just because your code doesnt cause runtime errors with the data that you are testing it with, doesnt mean that it isnt unsound from a typing perspective. Normally every interaction should be proveably valid just from the involved signatures. And i am using a type checker to actually ensure this.

But i guess that is not fully possible anyway with the setup of pandas.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I agree. We've chosen a middle ground with the stubs that supports the common ways people use pandas, but doesn't necessarily support all possible best typing practices.



def test_case_when() -> None:
Expand Down
Loading