Repository: yukinarit/pyserde
Pull request #32 (2020-09-25 13:54:09)

feat: Enum support
* You can serialize/deserialize a class including Enum/IntEnum
* Nested Enum is not yet supported


Patch:
diff --git a/README.md b/README.md
index 457f581..3da3f41 100644
--- a/README.md
+++ b/README.md
@@ -135,6 +135,7 @@ That's it! pyserde offers many more features. If you're interested, please read
 * Containers (List, Tuple, Dict)
 * [`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional)
 * User defined class with [`@dataclass`](https://docs.python.org/3/library/dataclasses.html)
+* [Enum](https://docs.python.org/3/library/enum.html#enum.Enum) and [IntEnum](https://docs.python.org/3/library/enum.html#enum.IntEnum)
 * More types
     * [`pathlib.Path`](https://docs.python.org/3/library/pathlib.html)
     * [`decimal.Decimal`](https://docs.python.org/3/library/decimal.html)
diff --git a/examples/basics.py b/examples/basics.py
index 7ab639b..b5b1f1c 100644
--- a/examples/basics.py
+++ b/examples/basics.py
@@ -2,6 +2,7 @@ from dataclasses import dataclass
 from serde import deserialize, serialize
 from serde.json import from_json, to_json
 
+
 # Mark the class serializable/deserializable.
 @deserialize
 @serialize
diff --git a/examples/enum34.py b/examples/enum34.py
new file mode 100644
index 0000000..a4c5f3a
--- /dev/null
+++ b/examples/enum34.py
@@ -0,0 +1,44 @@
+import enum
+from dataclasses import dataclass
+from serde import serialize, deserialize
+from serde.json import to_json, from_json
+
+
+class Nested(enum.Enum):
+    S = 'foo'
+
+
+class E(enum.Enum):
+    S = 'foo'
+    F = 10.0
+    B = True
+    N = Nested.S
+
+
+class IE(enum.IntEnum):
+    V0 = enum.auto()
+    V1 = enum.auto()
+    V2 = enum.auto()
+
+
+@deserialize
+@serialize
+@dataclass
+class Foo:
+    v0: IE
+    v1: IE = IE.V1  # Default enum value.
+    v2: E = E.S
+    # v3: E = E.N  # Sorry nested is not yet supported.
+
+
+if __name__ == "__main__":
+    f = Foo(IE.V0)
+    s = to_json(f)
+    print(s)
+
+    f = from_json(Foo, s)
+    print(f)
+    s = to_json(f)
+
+    # Non Enum values cause type mismatch error.
+    # s = to_json(Foo(0))
diff --git a/serde/compat.py b/serde/compat.py
index 5e36c2d..f935db3 100644
--- a/serde/compat.py
+++ b/serde/compat.py
@@ -1,6 +1,7 @@
 """
 Module for compatibility.
 """
+import enum
 from dataclasses import fields, is_dataclass
 from itertools import zip_longest
 from typing import Iterator, List, Optional, Tuple, Type, TypeVar, Union
@@ -160,6 +161,16 @@ def is_none(typ: Type) -> bool:
 PRIMITIVES = [int, float, bool, str]
 
 
+def is_enum(typ: Type) -> bool:
+    """
+    Test if the type is `enum.Enum`.
+    """
+    try:
+        return issubclass(typ, enum.Enum)
+    except TypeError:
+        return isinstance(typ, enum.Enum)
+
+
 def is_primitive(typ: Type) -> bool:
     """
     Test if the type is primitive.
diff --git a/serde/de.py b/serde/de.py
index f93d19d..db043d7 100644
--- a/serde/de.py
+++ b/serde/de.py
@@ -21,7 +21,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Type
 
 import jinja2
 
-from .compat import is_dict, is_list, is_opt, is_primitive, is_tuple, is_union, iter_types, type_args
+from .compat import is_dict, is_enum, is_list, is_opt, is_primitive, is_tuple, is_union, iter_types, type_args
 from .core import FROM_DICT, FROM_ITER, HIDDEN_NAME, SETTINGS, Field, Hidden, SerdeError, T, conv, fields, gen
 from .more_types import deserialize as custom
 
@@ -30,7 +30,7 @@ __all__: List = ['deserialize', 'is_deserializable', 'Deserializer', 'from_dict'
 Custom = Optional[Callable[['DeField', Any], Any]]
 
 
-def deserialize(_cls=None, rename_all: Optional[str] = None) -> Type:
+def deserialize(_cls=None, rename_all: Optional[str] = None):
     """
     `deserialize` decorator. A dataclass with this decorator can be deserialized
     into an object from various data format such as JSON and MsgPack.
@@ -68,6 +68,12 @@ def deserialize(_cls=None, rename_all: Optional[str] = None) -> Type:
         if not hasattr(cls, HIDDEN_NAME):
             setattr(cls, HIDDEN_NAME, Hidden())
         g['__custom_deserializer__'] = custom
+
+        # If there is a field of Enum, imports the class of Enum into the generation scope.
+        for f in dataclass_fields(cls):
+            if is_enum(f.type):
+                g[f.type.__name__] = f.type
+
         cls = de_func(cls, FROM_ITER, render_from_iter(cls, custom), g)
         cls = de_func(cls, FROM_DICT, render_from_dict(cls, rename_all, custom), g)
         return cls
@@ -281,6 +287,8 @@ class Renderer:
             return self.dict(arg)
         elif is_tuple(arg.type):
             return self.tuple(arg)
+        elif is_enum(arg.type):
+            return self.enum(arg)
         elif any(f(arg.type) for f in (is_primitive, is_union)):
             return self.primitive(arg)
         else:
@@ -369,6 +377,9 @@ class Renderer:
         k, v = arg.get_kv()
         return f'{{{self.render(k)}: {self.render(v)} for k, v in {arg.data}.items()}}'
 
+    def enum(self, arg: DeField) -> str:
+        return f'{arg.type.__name__}({self.primitive(arg)})'
+
     def primitive(self, arg: DeField) -> str:
         """
         Render rvalue for primitives.
diff --git a/serde/more_types.py b/serde/more_types.py
index 11b8509..1122585 100644
--- a/serde/more_types.py
+++ b/serde/more_types.py
@@ -1,6 +1,7 @@
 """
 Additional type support such as Decimal.
 """
+import enum
 from dataclasses import Field
 from decimal import Decimal
 from pathlib import Path
@@ -20,6 +21,8 @@ def serialize(data: Any) -> Any:
         return str(data)
     elif isinstance(data, Path):
         return str(data)
+    elif isinstance(data, enum.IntEnum):
+        return data.value
     else:
         raise SerdeError(f'Unsupported type: {type(data)}')
 
diff --git a/serde/se.py b/serde/se.py
index 0f76900..cc56fb7 100644
--- a/serde/se.py
+++ b/serde/se.py
@@ -4,6 +4,7 @@ Defines classess and functions for `serialize` decorator.
 """
 import abc
 import copy  # noqa
+import enum
 import functools
 from dataclasses import Field as DataclassField
 from dataclasses import asdict as _asdict
@@ -15,7 +16,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Type
 
 import jinja2
 
-from .compat import is_dict, is_list, is_opt, is_primitive, is_tuple, is_union, type_args
+from .compat import is_dict, is_enum, is_list, is_opt, is_primitive, is_tuple, is_union, type_args
 from .core import HIDDEN_NAME, SE_NAME, SETTINGS, TO_DICT, TO_ITER, Field, Hidden, T, conv, fields, gen
 from .more_types import serialize as custom
 
@@ -70,7 +71,7 @@ def serialize(_cls=None, rename_all: Optional[str] = None) -> Type:
     """
 
     @functools.wraps(_cls)
-    def wrap(cls: Type[T]) -> Type[T]:
+    def wrap(cls: Type[T]):
         if not hasattr(cls, HIDDEN_NAME):
             setattr(cls, HIDDEN_NAME, Hidden())
 
@@ -313,6 +314,8 @@ class Renderer:
             return self.dict(arg)
         elif is_tuple(arg.type):
             return self.tuple(arg)
+        elif is_enum(arg.type):
+            return self.enum(arg)
         elif any(f(arg.type) for f in (is_primitive, is_union)):
             return self.primitive(arg)
         else:
@@ -361,6 +364,9 @@ class Renderer:
         varg.name = 'v'
         return f'{{{self.render(karg)}: {self.render(varg)} for k, v in {arg.varname}.items()}}'
 
+    def enum(self, arg: SeField) -> str:
+        return f'{arg.varname}.value'
+
     def primitive(self, arg: SeField) -> str:
         """
         Render rvalue for primitives.


Test patch:
diff --git a/tests/test_basics.py b/tests/test_basics.py
index 6453fb8..48a7893 100644
--- a/tests/test_basics.py
+++ b/tests/test_basics.py
@@ -77,92 +77,85 @@ def test_forward_declaration():
     assert h.bar.i == 10
 
 
-@pytest.mark.parametrize('se,de', (format_dict + format_tuple + format_json + format_msgpack))
+@pytest.mark.parametrize('se,de', all_formats)
 def test_intenum(se, de):
     class IE(enum.IntEnum):
         V0 = enum.auto()
         V1 = enum.auto()
-        V2 = enum.auto()
+        V2 = 10
 
     @deserialize
     @serialize
     @dataclass
     class Foo:
-        ie0: IE
-        ie1: IE
-        ie2: IE
+        v0: IE
+        v1: IE
+        v2: IE = IE.V2
 
-    f = Foo(1, 2, 3)
-    assert f == de(Foo, se(f))
+    f = Foo(IE.V0, IE.V1)
+    ff = de(Foo, se(f))
+    assert f == ff
+    assert isinstance(ff.v0, enum.Enum)
+    assert isinstance(ff.v1, enum.Enum)
+    assert isinstance(ff.v2, enum.Enum)
 
-    f = Foo(IE.V0, IE.V1, IE.V2)
-    assert f == de(Foo, se(f))
+    # int literal doesn't work because of type mismatch.
+    f = Foo(1, 2, 3)
+    pytest.raises(Exception, lambda: de(Foo, se(f)))
 
 
-@pytest.mark.parametrize('se,de', format_yaml)
-def test_intenum_yaml(se, de):
-    class IE(enum.IntEnum):
-        V0 = enum.auto()
-        V1 = enum.auto()
-        V2 = enum.auto()
+@pytest.mark.parametrize('se,de', all_formats)
+def test_enum(se, de):
+    class E(enum.Enum):
+        V0 = 1
+        V1 = 'foo'
+        V2 = 10.0
+        V3 = True
 
     @deserialize
     @serialize
     @dataclass
     class Foo:
-        ie0: IE
-        ie1: IE
-        ie2: IE
+        v0: E
+        v1: E
+        v2: E
+        v3: E = E.V3
 
-    f = Foo(1, 2, 3)
-    assert f == de(Foo, se(f))
+    f = Foo(E.V0, E.V1, E.V2)
+    ff = de(Foo, se(f))
+    assert isinstance(ff.v0, enum.Enum)
+    assert isinstance(ff.v1, enum.Enum)
+    assert isinstance(ff.v2, enum.Enum)
+    assert isinstance(ff.v3, enum.Enum)
 
-    # yaml library doesn't serialize IntEnum.
-    f = Foo(IE.V0, IE.V1, IE.V2)
+    # literal values don't work because of type mismatch.
+    f = Foo(1, 'foo', 10.0, True)
     pytest.raises(Exception, lambda: de(Foo, se(f)))
 
 
-@pytest.mark.parametrize('se,de', format_toml)
-def test_intenum_toml(se, de):
-    class IE(enum.IntEnum):
+@pytest.mark.parametrize('se,de', all_formats)
+def test_enum_nested(se, de):
+    class Nested(enum.IntEnum):
         V0 = enum.auto()
         V1 = enum.auto()
         V2 = enum.auto()
 
-    @deserialize
-    @serialize
-    @dataclass
-    class Foo:
-        ie0: IE
-        ie1: IE
-        ie2: IE
-
-    f = Foo(1, 2, 3)
-    assert f == de(Foo, se(f))
-
-    # yaml library doesn't serialize IntEnum as integer but str.
-    f = Foo(IE.V0, IE.V1, IE.V2)
-    assert f != de(Foo, se(f))
-
-
-@pytest.mark.parametrize('se,de', all_formats)
-def test_enum(se, de):
     class E(enum.Enum):
-        V0 = 'v0'
-        V1 = 'v1'
-        V2 = 'v2'
+        V0 = Nested.V0
+        V1 = Nested.V1
+        V2 = Nested.V2
 
     @deserialize
     @serialize
     @dataclass
     class Foo:
-        ie0: E
-        ie1: E
-        ie2: E
+        v0: E
+        v1: E
+        v2: Nested = Nested.V0
 
-    # NOTE: Only IntEnum is supported now.
-    f = Foo(E.V0, E.V1, E.V2)
-    pytest.raises(Exception, lambda: de(Foo, se(f)))
+    # Nested enum don't work for yaml.
+    Foo(E.V0, E.V1)
+    # ff = de(Foo, se(f))
 
 
 @pytest.mark.parametrize('se,de', all_formats)