Skip to content

Commit

Permalink
Refine iterrows and itertuples return types (#822)
Browse files Browse the repository at this point in the history
refine iterrows and itertuples return types

* Added first test on iterrows.
* Updated iterrows return type to Iterable[tuple[Hashable, Series]].
* Added return type Iterable[tuple[Any, ...]] to itertuples.
  • Loading branch information
lebaudantoine authored Dec 1, 2023
1 parent 536db83 commit 2aadf0e
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 2 deletions.
6 changes: 4 additions & 2 deletions pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -247,8 +247,10 @@ class DataFrame(NDFrame, OpsMixin):
@property
def style(self) -> Styler: ...
def items(self) -> Iterable[tuple[Hashable, Series]]: ...
def iterrows(self) -> Iterable[tuple[Label, Series]]: ...
def itertuples(self, index: _bool = ..., name: _str | None = ...): ...
def iterrows(self) -> Iterable[tuple[Hashable, Series]]: ...
def itertuples(
self, index: _bool = ..., name: _str | None = ...
) -> Iterable[tuple[Any, ...]]: ...
def __len__(self) -> int: ...
@overload
def dot(self, other: DataFrame | ArrayLike) -> DataFrame: ...
Expand Down
6 changes: 6 additions & 0 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
getSeriesData,
)
from pandas.core.resample import Resampler # noqa: F401
from pandas.core.series import Series
import pytest
from typing_extensions import assert_type
import xarray as xr
Expand Down Expand Up @@ -403,6 +404,11 @@ def test_types_median() -> None:
s3: pd.Series = df.median(axis=1, skipna=True, numeric_only=False)


def test_types_iterrows() -> None:
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
res1: Iterable[tuple[Hashable, Series]] = df.iterrows()


def test_types_itertuples() -> None:
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
res1: Iterable[tuple[Any, ...]] = df.itertuples()
Expand Down

0 comments on commit 2aadf0e

Please sign in to comment.