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.

Asked By: Ondra Zahradnik

||

Answers:

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
Answered By: Martijn Pieters

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.

Answered By: Ethan Furman

I had the same issue in python3 and I found that the incompatibility between enum.Enum and pickle was because I was using enum.auto().

The solution was to put explicit fixed value on creation of my enum items and it solved the issue.

The reason is that enum.auto() generate an empty object which is unique, but it is identified only by its internal id (or its hash it is as you want) and from an execution to another, those id and hash and different, so, on a new execution of my app, pickle.load is unable to find the stored item in my enum class because the ids of the enum.auto() has changed.

I think that it is exacly the same issue for python2

Example :

use :

class Root_type(enum.Enum):
    no_scan = 0
    no_scan = 1

instead of :

class Root_type(enum.Enum):
    no_scan = seq.auto()
    no_scan = seq.auto()
Answered By: herve-guerin
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.