Skip to content

Fix value_counts and apply. #401

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.

Sign up for GitHub

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

Merged
merged 4 commits into from
Oct 27, 2022
Merged

Fix value_counts and apply. #401

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5305cd0
Series.value_counts returns Series[int].
danielroseman Oct 25, 2022
35df38f
Series.apply callable result might not be hashable
danielroseman Oct 25, 2022
e8d6a77
More detailed typing for DataFrame.apply.
danielroseman Oct 26, 2022
b1d4b25
Add test for #393.
danielroseman Oct 27, 2022
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
26 changes: 22 additions & 4 deletions pandas-stubs/core/frame.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1088,14 +1088,32 @@ class DataFrame(NDFrame, OpsMixin):
**kwargs,
) -> DataFrame: ...
@overload
def apply(self, f: Callable) -> Series: ...
def apply(
self,
f: Callable[..., Series],
axis: AxisType = ...,
raw: _bool = ...,
result_type: Literal["expand", "reduce", "broadcast"] | None = ...,
args=...,
**kwargs,
) -> DataFrame: ...
@overload
def apply(
self,
f: Callable[..., Scalar],
axis: AxisType = ...,
raw: _bool = ...,
result_type: Literal["expand", "reduce"] | None = ...,
args=...,
**kwargs,
) -> Series: ...
@overload
def apply(
self,
f: Callable,
axis: AxisType,
f: Callable[..., Scalar],
result_type: Literal["broadcast"],
axis: AxisType = ...,
raw: _bool = ...,
result_type: _str | None = ...,
args=...,
**kwargs,
) -> DataFrame: ...
Expand Down
4 changes: 2 additions & 2 deletions pandas-stubs/core/series.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -685,7 +685,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
@overload
def apply(
self,
func: Callable[..., Hashable],
func: Callable[..., Scalar | Sequence | Mapping],
Copy link
Contributor

Choose a reason for hiding this comment

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

Don't we basically assume everywhere in this code base that the expectation is that dataframes and series only store hashable/scalar types, and NOT things like dicts, or arbitrary objects for the matter?

Copy link
Collaborator

Choose a reason for hiding this comment

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

Yes, but in this case, it doesn't matter, because he is defining a Callable that returns types, that apply will then use to return a Series.

convertDType: _bool = ...,
args: tuple = ...,
**kwds,
Expand Down Expand Up @@ -1193,7 +1193,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
ascending: _bool = ...,
bins: int | None = ...,
dropna: _bool = ...,
) -> Series[S1]: ...
) -> Series[int]: ...
def transpose(self, *args, **kwargs) -> Series[S1]: ...
@property
def T(self) -> Series[S1]: ...
Expand Down
25 changes: 22 additions & 3 deletions tests/test_frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -460,9 +460,28 @@ def test_types_unique() -> None:

def test_types_apply() -> None:
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
df.apply(lambda x: x**2)
df.apply(np.exp)
df.apply(str)

def returns_series(x: pd.Series) -> pd.Series:
return x**2

check(assert_type(df.apply(returns_series), pd.DataFrame), pd.DataFrame)

def returns_scalar(x: pd.Series) -> float:
return 2

check(assert_type(df.apply(returns_scalar), pd.Series), pd.Series)
check(
assert_type(df.apply(returns_scalar, result_type="broadcast"), pd.DataFrame),
pd.DataFrame,
)
check(assert_type(df.apply(np.exp), pd.DataFrame), pd.DataFrame)
check(assert_type(df.apply(str), pd.Series), pd.Series)

# GH 393
def gethead(s: pd.Series, y: int) -> pd.Series:
return s.head(y)

check(assert_type(df.apply(gethead, args=(4,)), pd.DataFrame), pd.DataFrame)


def test_types_applymap() -> None:
Expand Down
9 changes: 7 additions & 2 deletions tests/test_series.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,8 +367,8 @@ def test_types_idxmax() -> None:


def test_types_value_counts() -> None:
s = pd.Series([1, 2])
s.value_counts()
s = pd.Series(["a", "b"])
check(assert_type(s.value_counts(), "pd.Series[int]"), pd.Series, int)


def test_types_unique() -> None:
Expand Down Expand Up @@ -398,6 +398,11 @@ def retseries(x: float) -> float:

check(assert_type(s.apply(retseries).tolist(), list), list)

def retlist(x: float) -> list:
return [x]

check(assert_type(s.apply(retlist), pd.Series), pd.Series, list)

def get_depth(url: str) -> int:
return len(url)

Expand Down