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 16, 2022
1 parent 3ee4b40 commit 2379fbd
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 4 deletions.
12 changes: 9 additions & 3 deletions src/cattr/converters.py
Original file line number Diff line number Diff line change
Expand Up @@ -405,9 +405,15 @@ 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__)
if val not in vals:
enum_vals = {x.value: x for x in vals if isinstance(x, Enum)}
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
38 changes: 37 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,30 @@ def test_structure_literal_multiple(converter_cls):

converter = converter_cls()

class Foo(Enum):
FOO = 7
FOOFOO = 77

class Bar(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
) == ClassWithLiteral(4)
assert converter.structure(
{"literal_field": 5}, ClassWithLiteral
) == ClassWithLiteral(5)
assert converter.structure(
{"literal_field": 7}, ClassWithLiteral
) == ClassWithLiteral(Foo.FOO)
assert converter.structure(
{"literal_field": 88}, ClassWithLiteral
) == ClassWithLiteral(Bar.BARBAR)


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

0 comments on commit 2379fbd

Please sign in to comment.