Search code examples
pythonpandaspython-typing

Type hint Pandas DataFrameGroupBy

How should I type hint in Python a pandas DataFrameGroupBy object? Should I just use pd.DataFrame as for normal pandas dataframes? I didn't find any other solution atm

Solution

  • DataFrameGroupBy is a proper type in of itself. So if you're writing a function which must specifically take a DataFrameGroupBy instance:

    from pandas.core.groupby import DataFrameGroupBy

def my_function(dfgb: DataFrameGroupBy) -> None:
    """Do something with dfgb."""

    If you're looking for a more general polymorphic type, there are several possibilities:

    • pandas.core.groupby.GroupBy since DataFrameGroupBy inherits from GroupBy[DataFrame].
    • If you want to accept Series instances too, you could either union DataFrameGroupBy and SeriesGroupBy or you could use GroupBy[FrameOrSeries] (if you intend to always match the input type in your return value) or GroupBy[FrameOrSeriesUnion] if your output type doesn't reflect the input type. All of these types are in pandas.core.groupby.generic.
    • You could combine the above generics (and others) in many different ways to your liking.