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.
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
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.
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()
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.
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
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.
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()