diff --git a/pkgs/clan-cli/clan_cli/api/serde.py b/pkgs/clan-cli/clan_cli/api/serde.py index 04c6ae2d5..394f077e2 100644 --- a/pkgs/clan-cli/clan_cli/api/serde.py +++ b/pkgs/clan-cli/clan_cli/api/serde.py @@ -54,6 +54,35 @@ def sanitize_string(s: str) -> str: return s +def is_enum(obj: Any) -> bool: + """ + Safely checks if the object or one of its attributes is an Enum. + """ + # Check if the object itself is an Enum + if isinstance(obj, Enum): + return True + + # Check if the object has an 'enum' attribute and if it's an Enum + enum_attr = getattr(obj, "enum", None) + return isinstance(enum_attr, Enum) + + +def get_enum_value(obj: Any) -> Any: + """ + Safely checks if the object or one of its attributes is an Enum. + """ + # Check if the object itself is an Enum + value = getattr(obj, "value", None) + if value is None and obj.enum: + value = getattr(obj.enum, "value", None) + + if value is None: + error_msg = f"Cannot determine enum value for {obj}" + raise ValueError(error_msg) + + return dataclass_to_dict(value) + + def dataclass_to_dict(obj: Any, *, use_alias: bool = True) -> Any: def _to_dict(obj: Any) -> Any: """ @@ -62,6 +91,8 @@ def dataclass_to_dict(obj: Any, *, use_alias: bool = True) -> Any: It does NOT convert member functions. """ + if is_enum(obj): + return get_enum_value(obj) if is_dataclass(obj): return { # Use either the original name or name diff --git a/pkgs/clan-cli/tests/test_serializers.py b/pkgs/clan-cli/tests/test_serializers.py index 927c71ac7..b265720ca 100644 --- a/pkgs/clan-cli/tests/test_serializers.py +++ b/pkgs/clan-cli/tests/test_serializers.py @@ -122,3 +122,18 @@ def test_filters_null_fields() -> None: assert instance.home == "home" assert instance.work is None assert dataclass_to_dict(instance) == {"home": "home"} + + +def test_custom_enum() -> None: + from enum import Enum + + class CustomEnum(Enum): + FOO = "foo" + BAR = "bar" + + @dataclass + class Foo: + field: CustomEnum + + instance = Foo(field=CustomEnum.FOO) + assert dataclass_to_dict(instance) == {"field": "foo"}