diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 1f8a2648..25c47523 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -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: ... diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 85943841..625b2824 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -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]: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 7d4890ee..13f9738d 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -599,32 +599,75 @@ 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, ) + for t1, t2 in df.iterrows(): + check(assert_type(t1, Hashable), Hashable) + check(assert_type(t2, pd.Series), pd.Series) 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, ) + for t1 in df.itertuples(): + assert_type(t1, _PandasNamedTuple) + assert t1.__class__.__name__ == "Pandas" + assert isinstance(t1.Index, int) + assert isinstance(t1.col1, int) + assert isinstance(t1.col2, int) + for k in [0, 1, 2]: + assert isinstance(t1[k], int) + + for t1 in df.itertuples(name="FooBar"): + assert_type(t1, _PandasNamedTuple) + assert t1.__class__.__name__ == "FooBar" + assert isinstance(t1.Index, int) + assert isinstance(t1.col1, int) + assert isinstance(t1.col2, int) + for k in [0, 1, 2]: + assert isinstance(t1[k], int) + + +def test_types_items() -> None: + df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]}) + check( + assert_type(df.items(), Iterator[tuple[Hashable, pd.Series]]), + Iterator, + tuple, + ) + + for t1, t2 in df.items(): + check(assert_type(t1, Hashable), Hashable) + check(assert_type(t2, pd.Series), pd.Series) + + +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]}) diff --git a/tests/test_series.py b/tests/test_series.py index 9edc8e06..98cfa65d 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import ( + Hashable, Iterable, Iterator, Sequence, @@ -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)