hashable Enum (#1461)

This commit is contained in:
Kevin Le 2022-10-27 04:55:38 -07:00 committed by GitHub
parent 6969023491
commit ccdd35b354
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 30 additions and 1 deletions

View File

@ -12,6 +12,10 @@ def eq_enum(self, other):
return self.value is other return self.value is other
def hash_enum(self):
return hash(self.name)
EnumType = type(PyEnum) EnumType = type(PyEnum)
@ -22,7 +26,7 @@ class EnumOptions(BaseOptions):
class EnumMeta(SubclassWithMeta_Meta): class EnumMeta(SubclassWithMeta_Meta):
def __new__(cls, name_, bases, classdict, **options): def __new__(cls, name_, bases, classdict, **options):
enum_members = dict(classdict, __eq__=eq_enum) enum_members = dict(classdict, __eq__=eq_enum, __hash__=hash_enum)
# We remove the Meta attribute from the class to not collide # We remove the Meta attribute from the class to not collide
# with the enum values. # with the enum values.
enum_members.pop("Meta", None) enum_members.pop("Meta", None)

View File

@ -518,3 +518,28 @@ def test_mutation_enum_input_type():
assert result.data == {"createPaint": {"color": "RED"}} assert result.data == {"createPaint": {"color": "RED"}}
assert color_input_value == RGB.RED assert color_input_value == RGB.RED
def test_hashable_enum():
class RGB(Enum):
"""Available colors"""
RED = 1
GREEN = 2
BLUE = 3
color_map = {RGB.RED: "a", RGB.BLUE: "b", 1: "c"}
assert color_map[RGB.RED] == "a"
assert color_map[RGB.BLUE] == "b"
assert color_map[1] == "c"
def test_hashable_instance_creation_enum():
Episode = Enum("Episode", [("NEWHOPE", 4), ("EMPIRE", 5), ("JEDI", 6)])
trilogy_map = {Episode.NEWHOPE: "better", Episode.EMPIRE: "best", 5: "foo"}
assert trilogy_map[Episode.NEWHOPE] == "better"
assert trilogy_map[Episode.EMPIRE] == "best"
assert trilogy_map[5] == "foo"