Pandas: How to assign sum() or mean() to df.groupby inside a function?

Question

I'd like to wrap df.groupby(pd.TimeGrouper(freq='M')).sum() in a function so that I can assign sum(), mean() or count() as arguments in that function. I've asked a similar question earlier here, but I don't think I can use the same technique in this particular case.

Here is a snippet with reproducible input:

# Imports
import pandas as pd
import numpy as np

# Dataframe with 1 or zero
# 100 rows and 4 columns
# Indexed by dates
np.random.seed(12345678)
df = pd.DataFrame(np.random.randint(0,2,size=(100, 4)), columns=list('ABCD'))
datelist = pd.date_range(pd.datetime(2017, 1, 1).strftime('%Y-%m-%d'), periods=100).tolist()
df['dates'] = datelist 
df = df.set_index(['dates'])
df.index = pd.to_datetime(df.index)
print(df.head(10))

Which gives:

With this we can do:

df2 = df.groupby(pd.TimeGrouper(freq='M')).sum()
print(df2)

And get:

Or we can do:

df3 = df.groupby(pd.TimeGrouper(freq='M')).mean()
print(df3)

And get:

Here's part of the procedure wrapped into a function:

# My function
def function1(df):
    df = df.groupby(pd.TimeGrouper(freq='M')).sum()
    return df

# Function1 call
df4 = function1(df = df)
print(df4)

And that works just fine:

The problem occurs when I try to add sum() or mean() as an argument in Function2, like this:

# My function with sum() as an argument
def function2(df, fun):
    df = df.groupby(pd.TimeGrouper(freq='M')).fun
    return df

My first attempt raises a TypeError:

# Function2 test 1
df5 = function2(df = df, fun = sum())

My second attempt raises an attribute error:

# Function2 test 2
df6 = function2(df = df, fun = 'sum()')

Is it possible to make a few adjustments to this setup to get it working? (I tried another version with 'M' as an argument for freq, and that worked just fine). Or is this just not the way these things are done?

Thank you for any suggestions!

Here is the whole mess for an easy copy&paste:

#%%

# Imports
import pandas as pd
import numpy as np

# Dataframe with 1 or zero
# 100 rows across 4 columns
# Indexed by dates
np.random.seed(12345678)
df = pd.DataFrame(np.random.randint(0,2,size=(100, 4)), columns=list('ABCD'))
datelist = pd.date_range(pd.datetime(2017, 1, 1).strftime('%Y-%m-%d'), periods=100).tolist()
df['dates'] = datelist 
df = df.set_index(['dates'])
df.index = pd.to_datetime(df.index)
print(df.head(10))

# Calculate sum per month
df2 = df.groupby(pd.TimeGrouper(freq='M')).sum()
print(df2)

# Or calculate average per month
df3 = df.groupby(pd.TimeGrouper(freq='M')).mean()
print(df3)

# My function
def function1(df):
    df = df.groupby(pd.TimeGrouper(freq='M')).sum()
    return df

# Function1 test
df4 = function1(df = df)
print(df4)
# So far so good
#%%
# My function with sum() as argument
def function2(df, fun):
    print(fun)
    df = df.groupby(pd.TimeGrouper(freq='M')).fun
    return df

# Function2 test 1
# df5 = function2(df = df, fun = sum())

# Function2 test 2
# df6 = function2(df = df, fun = 'sum()')

# Function2 test 3
# df7 = function2(df = df, fun = sum)

Pandas: How to assign sum() or mean() to df.groupby inside a function?

2 Answers