Type Hint For NamedTuple Returned By Pandas DataFrame itertuples()

Viewed 271

ITERTUPLES is a nice way to iterate over a pandas DF and it returns a namedtuple.

import pandas as pd
import numpy as np

df = pd.DataFrame({'num_legs': [4, 2], 'num_wings': [0, 2]},index=['dog', 'hawk'])
for row in df.itertuples():
    print(type(row))
    print(row)
<class 'pandas.core.frame.Pandas'>
Pandas(Index='dog', num_legs=4, num_wings=0)
<class 'pandas.core.frame.Pandas'>
Pandas(Index='hawk', num_legs=2, num_wings=2)

What is a correct way if any to add type hints to the returned namedtuples ?

2 Answers

I don't think its possible, because your dataframe can have any arbitrary data type, and thus the tuples will have any arbitrary data type present in the dataframe. In the same way you can't use Python type hints to specify the column types of a DataFrame, you can't explicitly type those named tuples.

If you need the type information of the columns before going into your for loop, you can certainly use df.dtypes, which gives you a Series with the column types.

One posible solution, in case the column names and data types are fixed, is to declare explicitly the data structure of the df row as a NamedTuple:

import pandas as pd

Row = NamedTuple(
    "Row",
    [("num_legs", int), ("num_wings", int), ("index", str)],
)

df = pd.DataFrame(
    {"num_legs": [4, 2], "num_wings": [0, 2]}, index=["dog", "hawk"]
)
row: Row
for row in df.itertuples():
    row.num_legs
Related