Pickling wrapped partial functions

Question

I am trying to create a picklable decorator using partial functions. However, I keep getting pickling errors when trying to do that.

The first naïve example is as follows:

def decorator(func):
  def wrapper(**kwargs):
    return partial(func, **kwargs)
  return wrapper

@decorator
def decorated(x, y=1, z=2):
  return x+y+z

y5 = decorated(y=5)
pickle.dumps(y5)

Where partial is taken from functools.

A bit less naïve attempt involves adding @wraps one line above def wrapper. This doesn't help.

I'm not sure I understand how pickling really works.

Answer 1

The problem is in your decorator, not with partial. A partial object should pickle just fine:

>>> from pickle import *
>>> from functools import *
>>> f = partial(pow, 2)
>>> p = dumps(f)
>>> g = loads(p)
>>> g(5)
32

So, this issue with your code is in the decorator. It is not preserving the name of the original function. Try this instead:

import pickle
from functools import *

def decorator(func):
    def wrapper(**kwargs):
        return partial(func, **kwargs)
    return wrapper

def decorated(x, y=1, z=2):
    return x+y+z

dd = decorator(decorated)

y5 = dd(y=5)
pickle.dumps(y5)

The modification to use dd should allow the pickle logic to discover the underlying function by its name. That is how pickles work.

To see the function name in the pickle, look at the dumps output:

>>> print pickle.dumps(y5)
cfunctools
partial
p0
(c__main__
decorated
p1
tp2
Rp3
(g1
(t(dp4
S'y'
p5
I5
sNtp6
b.

The word "decorated" needs to be findable, equal to the underlying function, and not be hidden by the decorator. Remember, when functions get pickled, only their names get stored. The actual contents of the function aren't in the pickle.

There are some workarounds, but they aren't pretty. You can use __setstate__() to save both the function name and its source-code. Then add a __getstate__() method to restore the function by exec-ing its source.

Alternatively, you can extract the byte codes in the function object object and save those. Upon a restore, compile the code object and exec it.

In short, your goal of using a decorator with the @ notation is directly at odds with how function pickling works. In order to achieve your goal, you'll have to customize function pickling to have it save what the function does, not just its name.