瀏覽代碼

added serializer for enums (#3058)

wassaf shahzad 1 年之前
父節點
當前提交
9073a2781b
共有 2 個文件被更改,包括 29 次插入5 次删除
  1. 14 0
      reflex/utils/serializers.py
  2. 15 5
      tests/utils/test_serializers.py

+ 14 - 0
reflex/utils/serializers.py

@@ -6,6 +6,7 @@ import json
 import types as builtin_types
 import warnings
 from datetime import date, datetime, time, timedelta
+from enum import Enum
 from typing import Any, Callable, Dict, List, Set, Tuple, Type, Union, get_type_hints
 
 from reflex.base import Base
@@ -232,6 +233,19 @@ def serialize_datetime(dt: Union[date, datetime, time, timedelta]) -> str:
     return str(dt)
 
 
+@serializer
+def serialize_enum(en: Enum) -> str:
+    """Serialize a enum to a JSON string.
+
+    Args:
+        en: The enum to serialize.
+
+    Returns:
+         The serialized enum.
+    """
+    return en.value
+
+
 @serializer
 def serialize_color(color: Color) -> str:
     """Serialize a color.

+ 15 - 5
tests/utils/test_serializers.py

@@ -12,11 +12,7 @@ from reflex.vars import Var
 
 @pytest.mark.parametrize(
     "type_,expected",
-    [
-        (str, True),
-        (dict, True),
-        (Dict[int, int], True),
-    ],
+    [(str, True), (dict, True), (Dict[int, int], True), (Enum, True)],
 )
 def test_has_serializer(type_: Type, expected: bool):
     """Test that has_serializer returns the correct value.
@@ -46,6 +42,7 @@ def test_has_serializer(type_: Type, expected: bool):
         (int, serializers.serialize_primitive),
         (float, serializers.serialize_primitive),
         (bool, serializers.serialize_primitive),
+        (Enum, serializers.serialize_enum),
     ],
 )
 def test_get_serializer(type_: Type, expected: serializers.Serializer):
@@ -103,6 +100,13 @@ class StrEnum(str, Enum):
     BAR = "bar"
 
 
+class TestEnum(Enum):
+    """A lone enum class."""
+
+    FOO = "foo"
+    BAR = "bar"
+
+
 class EnumWithPrefix(Enum):
     """An enum with a serializer adding a prefix."""
 
@@ -145,6 +149,12 @@ class BaseSubclass(Base):
             {"key1": EnumWithPrefix.FOO, "key2": EnumWithPrefix.BAR},
             '{"key1": "prefix_foo", "key2": "prefix_bar"}',
         ),
+        (TestEnum.FOO, "foo"),
+        ([TestEnum.FOO, TestEnum.BAR], '["foo", "bar"]'),
+        (
+            {"key1": TestEnum.FOO, "key2": TestEnum.BAR},
+            '{"key1": "foo", "key2": "bar"}',
+        ),
         (
             BaseSubclass(ts=datetime.timedelta(1, 1, 1)),
             '{"ts": "1 day, 0:00:01.000001"}',