fix: allow enum as mapping key (#3256)

vyperlang · Feb 1, 2023 · ddb402b · ddb402b
1 parent f3c8cb9
commit ddb402b
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/tests/parser/types/test_enum.py b/tests/parser/types/test_enum.py
@@ -276,3 +276,20 @@ def get_enum_from_struct() -> Foobar:
     """
     c = get_contract_with_gas_estimation(code)
     assert c.get_enum_from_struct() == 2
+
+
+def test_mapping_with_enum(get_contract_with_gas_estimation):
+    code = """
+enum Foobar:
+    FOO
+    BAR
+
+fb: HashMap[Foobar, uint256]
+
+@external
+def get_key(f: Foobar, i: uint256) -> uint256:
+    self.fb[f] = i
+    return self.fb[f]
+    """
+    c = get_contract_with_gas_estimation(code)
+    assert c.get_key(1, 777) == 777
diff --git a/vyper/semantics/types/user.py b/vyper/semantics/types/user.py
@@ -41,6 +41,7 @@ class EnumT(_UserType):
     # enums, but not static arrays of enums
     _as_darray = True
     _is_prim_word = True
+    _as_hashmap_key = True
 
     def __init__(self, name: str, members: dict) -> None:
         if len(members.keys()) > 256: