diff --git a/src/serialite/_dispatcher.py b/src/serialite/_dispatcher.py index 5fa9027..06a39c7 100644 --- a/src/serialite/_dispatcher.py +++ b/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): diff --git a/src/serialite/_implementations/__init__.py b/src/serialite/_implementations/__init__.py index b7796f3..20bf3b9 100644 --- a/src/serialite/_implementations/__init__.py +++ b/src/serialite/_implementations/__init__.py @@ -5,6 +5,7 @@ OrderedDictSerializer, RawDictSerializer, ) +from ._enum import EnumSerializer, InvalidEnumValueError from ._float import FloatSerializer from ._integer import ( IntegerOutOfRangeError, diff --git a/src/serialite/_implementations/_enum.py b/src/serialite/_implementations/_enum.py new file mode 100644 index 0000000..e929011 --- /dev/null +++ b/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}" diff --git a/tests/implementations/test_enum.py b/tests/implementations/test_enum.py new file mode 100644 index 0000000..24a298c --- /dev/null +++ b/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'" diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py index 62dfc2f..4eb0295 100644 --- a/tests/test_dispatcher.py +++ b/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):