Skip to content

Commit a370cab

Browse files
authored
itertuples (#842)
* Fix itertuples * _PandasNamedTuple * 3.9 * adjust old test * return Scalar
1 parent 888c3cc commit a370cab

File tree

2 files changed

+39
-5
lines changed

2 files changed

+39
-5
lines changed

pandas-stubs/core/frame.pyi

+4-1
Original file line numberDiff line numberDiff line change
@@ -249,7 +249,7 @@ class DataFrame(NDFrame, OpsMixin):
249249
def iterrows(self) -> Iterable[tuple[Hashable, Series]]: ...
250250
def itertuples(
251251
self, index: _bool = ..., name: _str | None = ...
252-
) -> Iterable[tuple[Any, ...]]: ...
252+
) -> Iterable[_PandasNamedTuple]: ...
253253
def __len__(self) -> int: ...
254254
@overload
255255
def dot(self, other: DataFrame | ArrayLike) -> DataFrame: ...
@@ -2279,3 +2279,6 @@ class DataFrame(NDFrame, OpsMixin):
22792279
) -> Self: ...
22802280
def __truediv__(self, other: float | DataFrame | Series | Sequence) -> Self: ...
22812281
def __rtruediv__(self, other: float | DataFrame | Series | Sequence) -> Self: ...
2282+
2283+
class _PandasNamedTuple(tuple[Any, ...]):
2284+
def __getattr__(self, field: str) -> Scalar: ...

tests/test_frame.py

+35-4
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,10 @@
3333
from pandas.core.resample import Resampler # noqa: F401
3434
from pandas.core.series import Series
3535
import pytest
36-
from typing_extensions import assert_type
36+
from typing_extensions import (
37+
TypeAlias,
38+
assert_type,
39+
)
3740
import xarray as xr
3841

3942
from pandas._typing import Scalar
@@ -47,6 +50,11 @@
4750
from pandas.io.formats.style import Styler
4851
from pandas.io.parsers import TextFileReader
4952

53+
if TYPE_CHECKING:
54+
from pandas.core.frame import _PandasNamedTuple
55+
else:
56+
_PandasNamedTuple: TypeAlias = tuple
57+
5058
DF = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
5159

5260

@@ -439,9 +447,23 @@ def test_types_iterrows() -> None:
439447

440448
def test_types_itertuples() -> None:
441449
df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})
442-
res1: Iterable[tuple[Any, ...]] = df.itertuples()
443-
res2: Iterable[tuple[Any, ...]] = df.itertuples(index=False, name="Foobar")
444-
res3: Iterable[tuple[Any, ...]] = df.itertuples(index=False, name=None)
450+
check(
451+
assert_type(df.itertuples(), Iterable[_PandasNamedTuple]),
452+
Iterable,
453+
_PandasNamedTuple,
454+
)
455+
check(
456+
assert_type(
457+
df.itertuples(index=False, name="Foobar"), Iterable[_PandasNamedTuple]
458+
),
459+
Iterable,
460+
_PandasNamedTuple,
461+
)
462+
check(
463+
assert_type(df.itertuples(index=False, name=None), Iterable[_PandasNamedTuple]),
464+
Iterable,
465+
_PandasNamedTuple,
466+
)
445467

446468

447469
def test_types_sum() -> None:
@@ -2962,3 +2984,12 @@ def test_frame_setitem_na() -> None:
29622984
df["x"] = df["y"] + pd.Timedelta(days=3)
29632985
df.loc[ind, :] = pd.NaT
29642986
df.iloc[[0, 2], :] = pd.NaT
2987+
2988+
2989+
def test_itertuples() -> None:
2990+
# GH 822
2991+
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})
2992+
2993+
for item in df.itertuples():
2994+
check(assert_type(item, _PandasNamedTuple), tuple)
2995+
assert_type(item.a, Scalar)

0 commit comments

Comments
 (0)