Home > Software design >  Type hint Pandas DataFrameGroupBy
Type hint Pandas DataFrameGroupBy

Time:12-29

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

CodePudding user response:

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.
  • Related