How to pickle complex enum values in Python

Question

When I try to unpickle pickled complex enum instance, I always get "ValueError: BLUE is not a valid Colors".

Is there any way how to pickle and unpickle?

from pickle import loads, dumps
from enum import Enum


class ColorType(object):

    def __init__(self, counter, name):
        self.counter = counter
        self.name = name

    def __str__(self):
        return self.name


class Colors(Enum):
    GREEN = ColorType(1, 'GREEN')
    BLUE = ColorType(2, 'BLUE')


color = Colors.BLUE
print(color is loads(dumps(color)))

I am using Python 2.7.

Answer 1

Don't use a custom class as the enum values; there is no need here. Your specific example doesn't need a separate class at all, you could just use:

class Colors(Enum):
    GREEN = 1
    BLUE = 2

    def __str__(self):
        return self.name

    @property
    def counter(self):
        return self.value

This has better str() and .counter behaviour; your code requires str() to be applied to Color.<name>.value rather than directly to Color.<name>.

For other custom methods and attributes, put those directly on the Enum subclass and they'll be part of the enum members too. If you need more values per entry, set a tuple and pull that tuple apart in a __init__ method. The documentation has an excellent Planet example that illustrates this further.

Demo:

>>> Colors.BLUE
<Colors.BLUE: 2>
>>> Colors.BLUE.value
2
>>> Colors.BLUE.counter
2
>>> str(Colors.BLUE)
'BLUE'
>>> Colors.BLUE is loads(dumps(Colors.BLUE))
True

Answer 2

The issue here is basic equality:

>>> ColorType(2, 'BLUE') == ColorType(2, 'BLUE')
False

So when Colors is trying to find a match for the unpickled value of ColorType(2, 'BLUE') it is failing.

The solution is simple: add the __eq__ and __ne__ methods to `ColorType':

class ColorType(object):

    def __init__(self, counter, name):
        self.counter = counter
        self.name = name

    def __str__(self):
        return self.name

    def __eq__(self, other):
        return self.name == other.name and self.counter == other.counter

    def __ne__(self, other):
        # not needed in Python 3
        return self.name != other .name or self.counter != other.counter

NB I agree with @MartijnPieters that in most cases you should just add the needed functionality to the Enum itself.