Skip to content

Commit

Permalink
Fix handling of Enums in Literal types
Browse files Browse the repository at this point in the history
  • Loading branch information
brakhane committed Feb 15, 2022
1 parent 3ee4b40 commit 92285e4
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 4 deletions.
14 changes: 11 additions & 3 deletions src/cattr/converters.py
Original file line number Diff line number Diff line change
Expand Up @@ -405,9 +405,17 @@ def _structure_call(obj, cl):

@staticmethod
def _structure_literal(val, type):
if val not in type.__args__:
raise Exception(f"{val} not in literal {type}")
return val
vals = set(type.__args__)
enums = {x for x in vals if isinstance(x, Enum)}
literal_vals = vals.difference(enums)
if val not in literal_vals:
enum_vals = {x.value: x for x in enums}
if val not in enum_vals:
raise Exception(f"{val} not in literal {type}")
else:
return enum_vals[val]
else:
return val

# Attrs classes.

Expand Down
41 changes: 40 additions & 1 deletion tests/test_structure_attrs.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -172,16 +194,33 @@ def test_structure_literal_multiple(converter_cls):

converter = converter_cls()

class Foo(Enum):
FOO = 1
FOOFOO = 2

class Bar(Enum):
BAR = 8
BARBAR = 9

@define
class ClassWithLiteral:
literal_field: Literal[4, 5] = 4
literal_field: Literal[4, 5, Literal[6], Foo.FOO, Bar.BARBAR] = 4

assert converter.structure(
{"literal_field": 4}, ClassWithLiteral
) == ClassWithLiteral(4)
assert converter.structure(
{"literal_field": 5}, ClassWithLiteral
) == ClassWithLiteral(5)
assert converter.structure(
{"literal_field": 6}, ClassWithLiteral
) == ClassWithLiteral(6)
assert converter.structure(
{"literal_field": 1}, ClassWithLiteral
) == ClassWithLiteral(Foo.FOO)
assert converter.structure(
{"literal_field": 9}, ClassWithLiteral
) == ClassWithLiteral(Bar.BARBAR)


@pytest.mark.skipif(is_py37, reason="Not supported on 3.7")
Expand Down

0 comments on commit 92285e4

Please sign in to comment.