diff --git a/platypush/schemas/__init__.py b/platypush/schemas/__init__.py index ce157319..6d12ce03 100644 --- a/platypush/schemas/__init__.py +++ b/platypush/schemas/__init__.py @@ -1,5 +1,6 @@ from datetime import datetime, date -from typing import Optional, Union +from enum import Enum +from typing import Any, Optional, Type, Union from dateutil.parser import isoparse from dateutil.tz import tzutc @@ -66,6 +67,23 @@ class Date(Function): # lgtm [py/missing-call-to-init] return date.fromtimestamp(dt.timestamp()) +class EnumField(Function): + """ + Field that maps enum values. + """ + + # pylint: disable=redefined-builtin + def __init__(self, *args, type: Type[Enum], **kwargs): + super().__init__(*args, **kwargs) + self.type = type + + def _serialize(self, value: Optional[Enum], *_, **__) -> Optional[Any]: + return value.value if value is not None else None + + def _deserialize(self, value: Optional[Any], *_, **__) -> Optional[Any]: + return self.type(value) if value is not None else None + + def normalize_datetime( dt: Optional[Union[str, date, datetime]] ) -> Optional[Union[date, datetime]]: