Change default arguments of function in python

Question

I have a function that converts a numpy array to a array containing True or False based on a condition then groups the True or False entries that are adjacent to one another and calculates the length of each group. This is to determine the length of dry spells or wet spells in a given month of precipitation data.

This is the function:

import itertools
def spell(X, kind='wet', how='mean', threshold=0.5): 

    if kind=='wet':
        condition = X>threshold
    else:
        condition = X<=threshold

    length = [sum(1 if x==True else nan for x in group) for key,group in itertools.groupby(condition) if key]

    if not length: 
        res = 0
    elif how=='mean': 
        res = np.mean(length)
    else:
        res = np.max(length)

    return res

So basically there is the option to determine the mean length or maximum length of wet or dry spells given a numpy array of precipitation data with the default parameters set to the mean length of wet spells.

I use this function with pandas to apply it to each month of a historical record:

#Create example dataframe
np.random.seed(1324)
idx = pd.DatetimeIndex(start='1960-01-01', periods=100, freq='d')
values = np.random.random(100)
df = pd.DataFrame(values, index=idx)

#Apply function
df.resample('M', how=spell)

and what I get is:

0
1960-01-31  1.555556
1960-02-29  1.500000
1960-03-31  1.777778
1960-04-30  6.000000

Which is perfect, however I want to be able to change the default values of this function somewhat on the fly so that I can use it's other options with df.resample(). I've looked into functools.partial() however this is only a solution for cases where the input arguments are explicitly set ie. spell(kind='dry', how='max', threshold=0.7). Is there a way to change the default arguments of the function in a way they will not need to be explicitly set afterwords so that I can used it with df.resample()?

Answer 1

The default values for a function are stored in that function's func_defaults attribute, which is a tuple of values which pair up with the trailing elements of the function's func_code.co_varnames tuple. For example:

>>> def foo(x, y=5):
...    return x, y
...
>>> foo(10)
(10, 5)
>>> foo.func_code.co_varnames
('x', 'y')
>>> foo.func_defaults
(5,)
>>> foo.func_defaults = (7,)
>>> foo(10)
(10, 7)

You can even give a parameter a default value after the fact:

>>> foo.func_defaults = (2, 3)
>>> foo()
(2, 3)

---

Warning: I had thought to (ab)use the mock library to allow temporarily overriding the function defaults, in a similar manner to a recent answer of mine. However, it seems to leave the defaults set to None afterward, which means either there is a bug in (or I misunderstand the behavior of) mock, or that messing with functions like this is a little dangerous.

def foo(x=5):
    return x

assert foo() == 5
with mock.patch.object(foo, 'func_defaults', (10,)):
    assert foo() == 10

assert foo() == 5  # Oops; I'm observing foo.func_defaults to be None now

Manually saving and restoring the defaults seems to work fine, though, as you might expect.

orig_defaults = foo.func_defaults
foo.func_defaults = (10,)
assert foo() == 10
foo.func_defaults = orig_defaults
assert foo() == 5

Answer 2

This sounds like a job for a function wrapper!

def spellwrapper(newkind, newhow, newthreshold):
    def wrapped_spell_func(X):
        spell(X, kind=newkind, how=newhow, threshold=newthreshold)
    return wrapped_spell_func

You would call this function with

new_spell_func = spellwrapper(newkind, newhow, newthreshold)

And it would return a wrapped version of the spell function that uses your new arguments as "defaults" instead of the ones created at the function definition. Then you would use

df.resample('M', how=new_spell_func)