From aca5f3c01c2714b6fbcd62e78dc7261a7f4f5fdf Mon Sep 17 00:00:00 2001 From: Lebaud Antoine Date: Fri, 1 Dec 2023 14:34:28 +0100 Subject: [PATCH] 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. --- pandas-stubs/core/frame.pyi | 6 ++++-- tests/test_frame.py | 6 ++++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 78ccbfb8e..6c9a11b6c 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -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: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index cb7d373f5..9bc32ddba 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -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 @@ -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()