Add support for enum type annotations

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 - Fixed serialization of optional nested dataclasses with a value of `None`.
 - Fixed preservation of comments on nested dataclass attributes.
+- Added support for using `enum.Enum` subclasses as type annotations.
 
 # 0.8.1 (2020-03-30)
 

datafiles/converters/__init__.py

@@ -1,4 +1,5 @@
 import dataclasses
+from enum import Enum
 from inspect import isclass
 from typing import Any, Dict, Optional, Union
 
@@ -9,6 +10,7 @@
 from ._bases import Converter
 from .builtins import Boolean, Float, Integer, String
 from .containers import Dataclass, Dictionary, List
+from .enumerations import Enumeration
 from .extensions import *  # pylint: disable=unused-wildcard-import
 
 
@@ -100,4 +102,7 @@ def map_type(cls, *, name: str = '', item_cls: Optional[type] = None):
         log.debug(f'Mapped {cls!r} to existing converter (itself)')
         return cls
 
+    if issubclass(cls, Enum):
+        return Enumeration.of_type(cls)
+
     raise TypeError(f'Could not map type: {cls}')

datafiles/converters/enumerations.py

@@ -0,0 +1,23 @@
+# pylint: disable=unused-argument,not-callable
+
+from ._bases import Converter
+
+
+class Enumeration(Converter):
+
+    ENUM: type = None  # type: ignore
+
+    @classmethod
+    def of_type(cls, enum: type):
+        name = f'{enum.__name__}Converter'
+        bases = (cls,)
+        attributes = {'ENUM': enum}
+        return type(name, bases, attributes)
+
+    @classmethod
+    def to_python_value(cls, deserialized_data, *, target_object=None):
+        return cls.ENUM(deserialized_data)
+
+    @classmethod
+    def to_preserialization_data(cls, python_value, *, default_to_skip=None):
+        return python_value.value

datafiles/tests/test_converters.py

@@ -1,6 +1,7 @@
 # pylint: disable=unused-variable
 
 from dataclasses import dataclass
+from enum import Enum
 from typing import ByteString, Dict, List, Optional
 
 import pytest
@@ -29,12 +30,17 @@ class MyCustomString:
     pass
 
 
+class Color(Enum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+
 IntegerList = converters.List.subclass(converters.Integer)
 StringList = converters.List.subclass(converters.String)
 MyDict = converters.Dictionary.subclass(converters.String, converters.Integer)
 MyDataclassConverter = converters.map_type(MyDataclass)
 MyDataclassConverterList = converters.map_type(List[MyDataclass])
-MyNestedDataclassConverter = converters.map_type(MyNestedDataclass)
 
 
 def describe_map_type():
@@ -68,6 +74,10 @@ def it_handles_dataclasses(expect):
             'flag': converters.Boolean,
         }
 
+    def it_handles_enums(expect):
+        converter = converters.map_type(Color)
+        expect(converter.__name__) == 'ColorConverter'
+
     def it_handles_optionals(expect):
         converter = converters.map_type(Optional[str])
         expect(converter.__name__) == 'OptionalString'
@@ -162,6 +172,12 @@ def when_text(expect):
             expect(convert("Hello, world!")) == "Hello, world!"
             expect(convert("Line 1\nLine 2\n")) == "Line 1\nLine 2\n"
 
+        def when_enum(expect):
+            convert = converters.map_type(Color).to_python_value
+            expect(convert(1)).is_(Color.RED)
+            with expect.raises(ValueError):
+                convert(42)
+
         def when_invalid(expect):
             message = "invalid literal for int() with base 10: 'a'"
             with expect.raises(ValueError, message):
@@ -240,6 +256,10 @@ def when_text(expect):
             expect(convert("Line 1\nLine 2")) == "Line 1\nLine 2\n"
             expect(convert("Line 1\nLine 2")).isinstance(LiteralScalarString)
 
+        def when_enum(expect):
+            convert = converters.map_type(Color).to_preserialization_data
+            expect(convert(Color.RED)) == 1
+
         def when_invalid(expect):
             message = "invalid literal for int() with base 10: 'a'"
             with expect.raises(ValueError, message):

docs/types/builtins.md

@@ -8,36 +8,53 @@ from typing import Optional
 
 ## Booleans
 
-| Type Annotation | Python Value | YAML Content |
-| --- | --- | --- |
-| `foobar: bool` | `foobar = True` | `foobar: true` |
-| `foobar: bool` | `foobar = False` | `foobar: false` |
-| `foobar: bool` | `foobar = None` | `foobar: false` |
-| `foobar: Optional[bool]` | `foobar = False` | `foobar: ` |
+| Type Annotation          | Python Value     | YAML Content    |
+| ------------------------ | ---------------- | --------------- |
+| `foobar: bool`           | `foobar = True`  | `foobar: true`  |
+| `foobar: bool`           | `foobar = False` | `foobar: false` |
+| `foobar: bool`           | `foobar = None`  | `foobar: false` |
+| `foobar: Optional[bool]` | `foobar = False` | `foobar:`       |
 
 ## Integers
 
-| Type Annotation | Python Value | YAML Content |
-| --- | --- | --- |
-| `foobar: int` | `foobar = 42` | `foobar: 42` |
-| `foobar: int` | `foobar = 1.23` | `foobar: 1` |
-| `foobar: int` | `foobar = None` | `foobar: 0` |
-| `foobar: Optional[int]` | `foobar = None` | `foobar: ` |
+| Type Annotation         | Python Value    | YAML Content |
+| ----------------------- | --------------- | ------------ |
+| `foobar: int`           | `foobar = 42`   | `foobar: 42` |
+| `foobar: int`           | `foobar = 1.23` | `foobar: 1`  |
+| `foobar: int`           | `foobar = None` | `foobar: 0`  |
+| `foobar: Optional[int]` | `foobar = None` | `foobar:`    |
 
 ## Floats
 
-| Type Annotation | Python Value | YAML Content |
-| --- | --- | --- |
-| `foobar: float` | `foobar = 1.23` | `foobar: 1.23` |
-| `foobar: float` | `foobar = 42` | `foobar: 42.0` |
-| `foobar: float` | `foobar = None` | `foobar: 0.0` |
-| `foobar: Optional[float]` | `foobar = None` | `foobar: ` |
+| Type Annotation           | Python Value    | YAML Content   |
+| ------------------------- | --------------- | -------------- |
+| `foobar: float`           | `foobar = 1.23` | `foobar: 1.23` |
+| `foobar: float`           | `foobar = 42`   | `foobar: 42.0` |
+| `foobar: float`           | `foobar = None` | `foobar: 0.0`  |
+| `foobar: Optional[float]` | `foobar = None` | `foobar:`      |
 
 ## Strings
 
-| Type Annotation | Python Value | YAML Content |
-| --- | --- | --- |
-| `foobar: str` | `foobar = "Hello, world!"` | `foobar: Hello, world!` |
-| `foobar: str` | `foobar = 42` | `foobar: '42'` |
-| `foobar: str` | `foobar = None` | `foobar: ''` |
-| `foobar: Optional[str]` | `foobar = None` | `foobar: ` |
+| Type Annotation         | Python Value               | YAML Content            |
+| ----------------------- | -------------------------- | ----------------------- |
+| `foobar: str`           | `foobar = "Hello, world!"` | `foobar: Hello, world!` |
+| `foobar: str`           | `foobar = 42`              | `foobar: '42'`          |
+| `foobar: str`           | `foobar = None`            | `foobar: ''`            |
+| `foobar: Optional[str]` | `foobar = None`            | `foobar:`               |
+
+## Enumerations
+
+Subclasses of `enum.Enum` can also be used as type annotations:
+
+```python
+from enum import Enum
+
+class Color:
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+```
+
+| Type Annotation | Python Value         | YAML Content |
+| --------------- | -------------------- | ------------ |
+| `color: Color`  | `color = Color.BLUE` | `color: 3`   |

docs/types/extensions.md

@@ -12,12 +12,12 @@ from typing import Optional
 from datafiles.converters import Number
 ```
 
-| Type Annotation | Python Value | YAML Content |
-| --- | --- | --- |
-| `foobar: Number` | `foobar = 42` | `foobar: 42` |
-| `foobar: Number` | `foobar = 1.23` | `foobar: 1.23` |
-| `foobar: Number` | `foobar = None` | `foobar: 0.0` |
-| `foobar: Optional[Number]` | `foobar = None` | `foobar: ` |
+| Type Annotation            | Python Value    | YAML Content   |
+| -------------------------- | --------------- | -------------- |
+| `foobar: Number`           | `foobar = 42`   | `foobar: 42`   |
+| `foobar: Number`           | `foobar = 1.23` | `foobar: 1.23` |
+| `foobar: Number`           | `foobar = None` | `foobar: 0.0`  |
+| `foobar: Optional[Number]` | `foobar = None` | `foobar:`      |
 
 ## Text
 
@@ -29,9 +29,9 @@ from typing import Optional
 from datafiles.converters import Text
 ```
 
-| Type Annotation | Python Value | YAML Content |
-| --- | --- | --- |
-| `foobar: Text` | `foobar = "Hello, world!"` | `foobar: Hello, world!` |
-| `foobar: Text` | `foobar = "First\nSecond\n"` | `foobar: | `<br>&nbsp;&nbsp;&nbsp;&nbsp;`First`<br>&nbsp;&nbsp;&nbsp;&nbsp;`Second` |
-| `foobar: Text` | `foobar = None` | `foobar: ''` |
-| `foobar: Optional[Text]` | `foobar = None` | `foobar: ` |
+| Type Annotation          | Python Value                 | YAML Content                                                                       |
+| ------------------------ | ---------------------------- | ---------------------------------------------------------------------------------- |
+| `foobar: Text`           | `foobar = "Hello, world!"`   | `foobar: Hello, world!`                                                            |
+| `foobar: Text`           | `foobar = "First\nSecond\n"` | `foobar: |`<br>&nbsp;&nbsp;&nbsp;&nbsp;`First`<br>&nbsp;&nbsp;&nbsp;&nbsp;`Second` |
+| `foobar: Text`           | `foobar = None`              | `foobar: ''`                                                                       |
+| `foobar: Optional[Text]` | `foobar = None`              | `foobar:`                                                                          |

mkdocs.yml

@@ -18,8 +18,8 @@ nav:
       - Mapper: api/mapper.md
   - Supported Types:
       - Builtins: types/builtins.md
-      - Extensions: types/extensions.md
       - Containers: types/containers.md
+      - Extended Types: types/extensions.md
       - Custom Types: types/custom.md
   - File Formats: formats.md
   - Utilities: utilities.md