Add support for Enum

src/serialite/_dispatcher.py

@@ -2,6 +2,7 @@
 
 from abc import get_cache_token
 from datetime import datetime
+from enum import Enum, IntEnum, StrEnum
 from pathlib import Path
 from types import GenericAlias, UnionType
 from typing import Any, Literal, NewType, TypeAliasType, Union, get_origin
@@ -245,6 +246,29 @@ def path_serializer(cls):
     return PathSerializer()
 
 
+@serializer.register(Enum)
+def enum_serializer(cls):
+    from ._implementations._enum import EnumSerializer
+
+    return EnumSerializer(cls)
+
+
+# int appears before Enum in IntEnum's MRO, so it would dispatch to int without this
+@serializer.register(IntEnum)
+def int_enum_serializer(cls):
+    from ._implementations._enum import EnumSerializer
+
+    return EnumSerializer(cls)
+
+
+# str appears before Enum in StrEnum's MRO, so it would dispatch to str without this
+@serializer.register(StrEnum)
+def str_enum_serializer(cls):
+    from ._implementations._enum import EnumSerializer
+
+    return EnumSerializer(cls)
+
+
 # Union disables subclassing so Optional cannot be used to dispatch
 # @serializer.register(Optional)
 def optional_serializer(cls):

src/serialite/_implementations/__init__.py

@@ -5,6 +5,7 @@
     OrderedDictSerializer,
     RawDictSerializer,
 )
+from ._enum import EnumSerializer, InvalidEnumValueError
 from ._float import FloatSerializer
 from ._integer import (
     IntegerOutOfRangeError,

src/serialite/_implementations/_enum.py

@@ -0,0 +1,85 @@
+__all__ = ["EnumSerializer", "InvalidEnumValueError"]
+
+from dataclasses import dataclass
+from enum import Enum, IntEnum, StrEnum
+from typing import Any, Literal
+
+from .._base import Serializer
+from .._decorators import serializable
+from .._errors import Errors
+from .._numeric_check import is_int
+from .._result import Failure, Result, Success
+from .._type_errors import ExpectedIntegerError, ExpectedStringError
+
+
+class EnumSerializer[E: Enum](Serializer[E]):
+    def __init__(self, enum_class: type[E], *, by: Literal["name", "value"] = "name"):
+        if by not in ["name", "value"]:
+            raise ValueError(f"Expected 'name' or 'value' for by, but got {by!r}")
+
+        self.enum_class = enum_class
+        self.by = by
+
+    def from_data(self, data) -> Result[E]:
+        if self.by == "name":
+            return self._from_data_by_name(data)
+        else:
+            return self._from_data_by_value(data)
+
+    def _from_data_by_name(self, data) -> Result[E]:
+        if not isinstance(data, str):
+            return Failure(Errors.one(ExpectedStringError(data)))
+
+        try:
+            return Success(self.enum_class[data])
+        except KeyError:
+            names = [m.name for m in self.enum_class]
+            err = InvalidEnumValueError(self.enum_class.__name__, names, data)
+            return Failure(Errors.one(err))
+
+    def _from_data_by_value(self, data) -> Result[E]:
+        if issubclass(self.enum_class, StrEnum):
+            if not isinstance(data, str):
+                return Failure(Errors.one(ExpectedStringError(data)))
+        elif issubclass(self.enum_class, IntEnum):
+            if not is_int(data):
+                return Failure(Errors.one(ExpectedIntegerError(data)))
+
+        try:
+            return Success(self.enum_class(data))
+        except ValueError:
+            values = [m.value for m in self.enum_class]
+            err = InvalidEnumValueError(self.enum_class.__name__, values, data)
+            return Failure(Errors.one(err))
+
+    def to_data(self, value: E):
+        if not isinstance(value, self.enum_class):
+            raise ValueError(f"Expected {self.enum_class.__name__}, got {type(value).__name__}")
+
+        if self.by == "name":
+            return value.name
+        else:
+            return value.value
+
+    def to_openapi_schema(self, force: bool = False):
+        if self.by == "name":
+            return {"type": "string", "enum": [m.name for m in self.enum_class]}
+        else:
+            values = [m.value for m in self.enum_class]
+            if all(isinstance(v, str) for v in values):
+                return {"type": "string", "enum": values}
+            elif all(is_int(v) for v in values):
+                return {"type": "integer", "enum": values}
+            else:
+                return {"enum": values}
+
+
+@serializable
+@dataclass(frozen=True, slots=True)
+class InvalidEnumValueError(Exception):
+    enum_name: str
+    values: list[Any]
+    actual: Any
+
+    def __str__(self) -> str:
+        return f"Expected one of {self.values!r} for {self.enum_name}, but got {self.actual!r}"

tests/implementations/test_enum.py

@@ -0,0 +1,170 @@
+from datetime import datetime
+from enum import Enum, IntEnum, StrEnum, auto
+
+import pytest
+
+from serialite import (
+    EnumSerializer,
+    Errors,
+    ExpectedIntegerError,
+    ExpectedStringError,
+    Failure,
+    InvalidEnumValueError,
+    Success,
+)
+
+
+class Color(Enum):
+    RED = "red"
+    GREEN = "green"
+
+
+class Priority(IntEnum):
+    LOW = 1
+    HIGH = auto()
+
+
+class Status(StrEnum):
+    ACTIVE = "active"
+    INACTIVE = auto()
+
+
+class AutoEnum(Enum):
+    A = auto()
+    B = auto()
+
+
+DATE = datetime(2024, 1, 1, 12, 0, 0)
+
+
+class MixedEnum(Enum):
+    STRING = "hello"
+    NUMBER = 42
+    DATE = DATE
+
+
+# By name (default)
+@pytest.mark.parametrize(
+    ("enum_class", "name", "member"),
+    [
+        (Color, "RED", Color.RED),
+        (Color, "GREEN", Color.GREEN),
+        (Priority, "LOW", Priority.LOW),
+        (Priority, "HIGH", Priority.HIGH),
+        (Status, "ACTIVE", Status.ACTIVE),
+        (Status, "INACTIVE", Status.INACTIVE),
+        (AutoEnum, "A", AutoEnum.A),
+        (AutoEnum, "B", AutoEnum.B),
+        (MixedEnum, "STRING", MixedEnum.STRING),
+        (MixedEnum, "DATE", MixedEnum.DATE),
+    ],
+)
+def test_by_name(enum_class, name, member):
+    s = EnumSerializer(enum_class)
+    assert s.from_data(name) == Success(member)
+    assert s.to_data(member) == name
+
+
+@pytest.mark.parametrize("data", [123, True, None, ["a"]])
+def test_by_name_rejects_non_string_data(data):
+    s = EnumSerializer(Color)
+    expected = ExpectedStringError(data)
+    assert s.from_data(data) == Failure(Errors.one(expected))
+
+
+def test_by_name_rejects_unknown_name():
+    s = EnumSerializer(Color)
+    expected = InvalidEnumValueError("Color", ["RED", "GREEN"], "BLUE")
+    assert s.from_data("BLUE") == Failure(Errors.one(expected))
+
+
+def test_by_name_to_data_raises_on_non_member():
+    s = EnumSerializer(Color)
+    with pytest.raises(ValueError):
+        s.to_data("RED")
+
+
+@pytest.mark.parametrize(
+    ("enum_class", "expected"),
+    [
+        (Color, {"type": "string", "enum": ["RED", "GREEN"]}),
+        (Priority, {"type": "string", "enum": ["LOW", "HIGH"]}),
+        (Status, {"type": "string", "enum": ["ACTIVE", "INACTIVE"]}),
+        (AutoEnum, {"type": "string", "enum": ["A", "B"]}),
+        (MixedEnum, {"type": "string", "enum": ["STRING", "NUMBER", "DATE"]}),
+    ],
+)
+def test_by_name_openapi_schema(enum_class, expected):
+    s = EnumSerializer(enum_class)
+    assert s.to_openapi_schema() == expected
+
+
+# By value
+@pytest.mark.parametrize(
+    ("enum_class", "value", "member"),
+    [
+        (Color, "red", Color.RED),
+        (Color, "green", Color.GREEN),
+        (Priority, 1, Priority.LOW),
+        (Priority, 2, Priority.HIGH),
+        (Status, "active", Status.ACTIVE),
+        (Status, "inactive", Status.INACTIVE),
+        (MixedEnum, "hello", MixedEnum.STRING),
+        (MixedEnum, 42, MixedEnum.NUMBER),
+        (MixedEnum, DATE, MixedEnum.DATE),
+    ],
+)
+def test_by_value(enum_class, value, member):
+    s = EnumSerializer(enum_class, by="value")
+    assert s.from_data(value) == Success(member)
+    assert s.to_data(member) == value
+
+
+@pytest.mark.parametrize(
+    ("enum_class", "data", "expected"),
+    [
+        (Priority, "1", ExpectedIntegerError("1")),
+        (Priority, True, ExpectedIntegerError(True)),
+        (Status, 123, ExpectedStringError(123)),
+    ],
+)
+def test_by_value_rejects_wrong_type(enum_class, data, expected):
+    s = EnumSerializer(enum_class, by="value")
+    assert s.from_data(data) == Failure(Errors.one(expected))
+
+
+def test_by_value_rejects_unknown_value():
+    s = EnumSerializer(Color, by="value")
+    expected = InvalidEnumValueError("Color", ["red", "green"], "blue")
+    assert s.from_data("blue") == Failure(Errors.one(expected))
+
+
+def test_by_value_to_data_raises_on_non_member():
+    s = EnumSerializer(Color, by="value")
+    with pytest.raises(ValueError):
+        s.to_data("red")
+
+
+@pytest.mark.parametrize(
+    ("enum_class", "expected"),
+    [
+        (Color, {"type": "string", "enum": ["red", "green"]}),
+        (Priority, {"type": "integer", "enum": [1, 2]}),
+        (Status, {"type": "string", "enum": ["active", "inactive"]}),
+        (MixedEnum, {"enum": ["hello", 42, DATE]}),
+    ],
+)
+def test_by_value_openapi_schema(enum_class, expected):
+    s = EnumSerializer(enum_class, by="value")
+    assert s.to_openapi_schema() == expected
+
+
+# Error
+def test_invalid_enum_value_error():
+    error = InvalidEnumValueError("Color", ["RED", "GREEN"], "BLUE")
+    assert error.to_data() == {
+        "enum_name": "Color",
+        "values": ["RED", "GREEN"],
+        "actual": "BLUE",
+    }
+    assert str(error) == "Expected one of ['RED', 'GREEN'] for Color, but got 'BLUE'"

tests/test_dispatcher.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from enum import Enum, IntEnum, StrEnum
 from typing import Any, Dict, List, Literal, NewType, Optional, Tuple, Union
 from uuid import UUID
 
@@ -7,6 +8,21 @@
 from serialite import Success, serializer
 
 
+class Color(Enum):
+    RED = "red"
+    GREEN = "green"
+
+
+class Priority(IntEnum):
+    LOW = 1
+    HIGH = 2
+
+
+class Status(StrEnum):
+    ACTIVE = "active"
+    INACTIVE = "inactive"
+
+
 @pytest.mark.parametrize(
     ("data_type", "data", "value"),
     [
@@ -27,6 +43,9 @@
         (tuple[int, str], [5, "a"], (5, "a")),
         (Dict[str, int], {"a": 11, "b": 22}, {"a": 11, "b": 22}),
         (dict[str, int], {"a": 11, "b": 22}, {"a": 11, "b": 22}),
+        (Color, "RED", Color.RED),
+        (Priority, "LOW", Priority.LOW),
+        (Status, "ACTIVE", Status.ACTIVE),
     ],
 )
 def test_dispatch(data_type, data, value):