[Python] How to pickle Enum with values of type functools.partial

Problem

Suppose we have a python Enum where values are of type functools.partial. How to pickle and unpickle a member of that enum ?

import pickle
from enum import Enum
from functools import partial


def function_a():
pass


class EnumOfPartials(Enum):
FUNCTION_A = partial(function_a)


if __name__ == "__main__":
with open("test.pkl", "wb") as f:
pickle.dump(EnumOfPartials.FUNCTION_A, f)

with open("test.pkl", "rb") as f:
pickle.load(f)


The code above tries to pickle and unpickle such an object. The pickle.load operations results in error:

ValueError: functools.partial(<function function_a at 0x7f5973e804a0>) is not a valid EnumOfPartials


Motivation

The object in itself is useful for configurations purposes: using hydra, I can have a parameter in a YAML that corresponds to the choice of a function in the enum. The reason for using partial is so that FUNCTION_A does not get interpreted as a method (see this question). Being able to pickle a member of this enum is desirable to be able to send it to another process.

Given my use case, an obvious workaround would be to have a dictionary of functions indexed by an enum, but I would prefer directly having the relevant value (the function) in the enum.

Note

I am using python 3.11.11.

Continue reading...