Corrected enum metaclass to fix pickle.dumps()

This commit is contained in:
sgrekov 2023-02-13 17:22:21 +02:00
parent 52143473ef
commit f5b5a14f3f
2 changed files with 30 additions and 1 deletions

View File

@ -0,0 +1,27 @@
import pickle
from ...types.enum import Enum
class MyEnum(Enum):
# is defined outside of test because pickle unable to dump class inside ot pytest function
A = 'a'
B = 1
def test_enums_pickling():
a = MyEnum.A
pickled = pickle.dumps(a)
restored = pickle.loads(pickled)
assert type(a) is type(restored)
assert a == restored
assert a.value == restored.value
assert a.name == restored.name
b = MyEnum.B
pickled = pickle.dumps(b)
restored = pickle.loads(pickled)
assert type(a) is type(restored)
assert b == restored
assert b.value == restored.value
assert b.name == restored.name

View File

@ -31,9 +31,11 @@ class EnumMeta(SubclassWithMeta_Meta):
# with the enum values. # with the enum values.
enum_members.pop("Meta", None) enum_members.pop("Meta", None)
enum = PyEnum(cls.__name__, enum_members) enum = PyEnum(cls.__name__, enum_members)
return SubclassWithMeta_Meta.__new__( obj = SubclassWithMeta_Meta.__new__(
cls, name_, bases, dict(classdict, __enum__=enum), **options cls, name_, bases, dict(classdict, __enum__=enum), **options
) )
globals()[name_] = obj.__enum__
return obj
def get(cls, value): def get(cls, value):
return cls._meta.enum(value) return cls._meta.enum(value)