Fix handling of Enums in Literal types

Author: brakhane

HISTORY.rst

@@ -11,10 +11,11 @@ History
   (`#218 <https://github.com/python-attrs/cattrs/issues/218>`_)
 * Fix structuring bare ``typing.Tuple`` on Pythons lower than 3.9.
   (`#218 <https://github.com/python-attrs/cattrs/issues/218>`_)
-
 * Fix a wrong ``AttributeError`` of an missing ``__parameters__`` attribute. This could happen 
   when inheriting certain generic classes – for example ``typing.*`` classes are affected. 
   (`#217 <https://github.com/python-attrs/cattrs/issues/217>`_)
+* Fix structuring of ``enum.Enum`` instances in ``typing.Literal`` types.
+  (`#231 <https://github.com/python-attrs/cattrs/pull/231>`_)
 
 1.10.0 (2022-01-04)
 -------------------

src/cattr/converters.py

@@ -86,6 +86,12 @@ def is_optional(typ):
     )
 
 
+def is_literal_containing_enums(typ):
+    return is_literal(typ) and any(
+        isinstance(val, Enum) for val in typ.__args__
+    )
+
+
 class Converter(object):
     """Converts between structured and unstructured data."""
 
@@ -155,7 +161,8 @@ def __init__(
                     lambda v, _: v,
                 ),
                 (is_generic_attrs, self._gen_structure_generic, True),
-                (is_literal, self._structure_literal),
+                (is_literal, self._structure_simple_literal),
+                (is_literal_containing_enums, self._structure_enum_literal),
                 (is_sequence, self._structure_list),
                 (is_mutable_set, self._structure_set),
                 (is_frozenset, self._structure_frozenset),
@@ -404,11 +411,21 @@ def _structure_call(obj, cl):
         return cl(obj)
 
     @staticmethod
-    def _structure_literal(val, type):
+    def _structure_simple_literal(val, type):
         if val not in type.__args__:
             raise Exception(f"{val} not in literal {type}")
         return val
 
+    @staticmethod
+    def _structure_enum_literal(val, type):
+        vals = {
+            (x.value if isinstance(x, Enum) else x): x for x in type.__args__
+        }
+        try:
+            return vals[val]
+        except KeyError:
+            raise Exception(f"{val} not in literal {type}") from None
+
     # Attrs classes.
 
     def structure_attrs_fromtuple(

tests/test_structure_attrs.py

@@ -1,4 +1,5 @@
 """Loading of attrs classes."""
+from enum import Enum
 from ipaddress import IPv4Address, IPv6Address, ip_address
 from typing import Union
 from unittest.mock import Mock
@@ -164,6 +165,27 @@ class ClassWithLiteral:
     ) == ClassWithLiteral(4)
 
 
+@pytest.mark.skipif(is_py37, reason="Not supported on 3.7")
+@pytest.mark.parametrize("converter_cls", [Converter, GenConverter])
+def test_structure_literal_enum(converter_cls):
+    """Structuring a class with a literal field works."""
+    from typing import Literal
+
+    converter = converter_cls()
+
+    class Foo(Enum):
+        FOO = 1
+        BAR = 2
+
+    @define
+    class ClassWithLiteral:
+        literal_field: Literal[Foo.FOO] = Foo.FOO
+
+    assert converter.structure(
+        {"literal_field": 1}, ClassWithLiteral
+    ) == ClassWithLiteral(Foo.FOO)
+
+
 @pytest.mark.skipif(is_py37, reason="Not supported on 3.7")
 @pytest.mark.parametrize("converter_cls", [Converter, GenConverter])
 def test_structure_literal_multiple(converter_cls):
@@ -172,9 +194,17 @@ def test_structure_literal_multiple(converter_cls):
 
     converter = converter_cls()
 
+    class Foo(Enum):
+        FOO = 7
+        FOOFOO = 77
+
+    class Bar(int, Enum):
+        BAR = 8
+        BARBAR = 88
+
     @define
     class ClassWithLiteral:
-        literal_field: Literal[4, 5] = 4
+        literal_field: Literal[4, 5, Foo.FOO, Bar.BARBAR] = 4
 
     assert converter.structure(
         {"literal_field": 4}, ClassWithLiteral
@@ -183,6 +213,14 @@ class ClassWithLiteral:
         {"literal_field": 5}, ClassWithLiteral
     ) == ClassWithLiteral(5)
 
+    assert converter.structure(
+        {"literal_field": 7}, ClassWithLiteral
+    ) == ClassWithLiteral(Foo.FOO)
+
+    cwl = converter.structure({"literal_field": 88}, ClassWithLiteral)
+    assert cwl == ClassWithLiteral(Bar.BARBAR)
+    assert isinstance(cwl.literal_field, Bar)
+
 
 @pytest.mark.skipif(is_py37, reason="Not supported on 3.7")
 @pytest.mark.parametrize("converter_cls", [Converter, GenConverter])