Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
8 changes: 4 additions & 4 deletions pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -328,16 +328,16 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
def shape(self) -> tuple[int, int]: ...
@property
def style(self) -> Styler: ...
def items(self) -> Iterable[tuple[Hashable, Series]]: ...
def iterrows(self) -> Iterable[tuple[Hashable, Series]]: ...
def items(self) -> Iterator[tuple[Hashable, Series]]: ...
def iterrows(self) -> Iterator[tuple[Hashable, Series]]: ...
@overload
def itertuples(
self, index: _bool = ..., name: _str = ...
) -> Iterable[_PandasNamedTuple]: ...
) -> Iterator[_PandasNamedTuple]: ...
@overload
def itertuples(
self, index: _bool = ..., name: None = None
) -> Iterable[tuple[Any, ...]]: ...
) -> Iterator[tuple[Any, ...]]: ...
def __len__(self) -> int: ...
@overload
def dot(self, other: DataFrame | ArrayLike) -> Self: ...
Expand Down
2 changes: 1 addition & 1 deletion pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -615,7 +615,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
mode: Literal["w"] = ...,
) -> _str: ...
def to_xarray(self) -> xr.DataArray: ...
def items(self) -> Iterable[tuple[Hashable, S1]]: ...
def items(self) -> Iterator[tuple[Hashable, S1]]: ...
def keys(self) -> Index: ...
@overload
def to_dict(self, *, into: type[dict] = ...) -> dict[Any, S1]: ...
Expand Down
23 changes: 16 additions & 7 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -599,7 +599,7 @@ def test_types_median() -> None:
def test_types_iterrows() -> None:
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
check(
assert_type(df.iterrows(), "Iterable[tuple[Hashable, pd.Series]]"),
assert_type(df.iterrows(), "Iterator[tuple[Hashable, pd.Series]]"),
Iterable,
tuple,
)
Comment on lines 601 to 605
Copy link
Collaborator

Choose a reason for hiding this comment

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

Can we add a test of this form for iterrows() and itertuples():

for t1, t2 in df.iterrows():
    check(assert_type(t1, Hashable), Hashable)
    check(assert_type(t2, pd.Series), pd.Series)

Expand All @@ -608,24 +608,33 @@ def test_types_iterrows() -> None:
def test_types_itertuples() -> None:
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
check(
assert_type(df.itertuples(), Iterable[_PandasNamedTuple]),
Iterable,
assert_type(df.itertuples(), Iterator[_PandasNamedTuple]),
Iterator,
_PandasNamedTuple,
)
check(
assert_type(
df.itertuples(index=False, name="Foobar"), Iterable[_PandasNamedTuple]
df.itertuples(index=False, name="Foobar"), Iterator[_PandasNamedTuple]
),
Iterable,
Iterator,
_PandasNamedTuple,
)
check(
assert_type(df.itertuples(index=False, name=None), Iterable[tuple[Any, ...]]),
Iterable,
assert_type(df.itertuples(index=False, name=None), Iterator[tuple[Any, ...]]),
Iterator,
object,
)


def test_frame_iterator() -> None:
"""Test iterator methods for a dataframe GH1217."""
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})

check(assert_type(next(df.items()), tuple[Hashable, "pd.Series"]), tuple)
check(assert_type(next(df.iterrows()), tuple[Hashable, "pd.Series"]), tuple)
check(assert_type(next(df.itertuples()), _PandasNamedTuple), _PandasNamedTuple)


def test_types_sum() -> None:
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
df.sum()
Expand Down
7 changes: 7 additions & 0 deletions tests/test_series.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

from collections.abc import (
Hashable,
Iterable,
Iterator,
Sequence,
Expand Down Expand Up @@ -3797,3 +3798,9 @@ def foo(sf: pd.Series) -> None:

foo(s)
check(assert_type(s + pd.Series([1]), pd.Series), pd.Series)


def test_series_items() -> None:
s = pd.Series(data=[1, 2, 3, 4], index=["cow", "coal", "coalesce", ""])
check(assert_type(next(s.items()), tuple[Hashable, int]), tuple)
check(assert_type(s.items(), Iterator[tuple[Hashable, int]]), Iterator)