diff --git a/packages/polywrap-msgpack/polywrap_msgpack/sanitize.py b/packages/polywrap-msgpack/polywrap_msgpack/sanitize.py index 6c7407a2..2b67d136 100644 --- a/packages/polywrap-msgpack/polywrap_msgpack/sanitize.py +++ b/packages/polywrap-msgpack/polywrap_msgpack/sanitize.py @@ -2,6 +2,7 @@ python values into msgpack compatible values.""" from __future__ import annotations +from enum import IntEnum from typing import Any, Dict, List, Set, Tuple, cast from .extensions.generic_map import GenericMap @@ -42,6 +43,8 @@ def sanitize(value: Any) -> Any: ... ValueError: GenericMap key must be string, got 1 of type """ + if isinstance(value, IntEnum): + return value.value if isinstance(value, GenericMap): dictionary: Dict[Any, Any] = cast( GenericMap[Any, Any], value diff --git a/packages/polywrap-msgpack/pyproject.toml b/packages/polywrap-msgpack/pyproject.toml index 5c55f8ab..0be2dced 100644 --- a/packages/polywrap-msgpack/pyproject.toml +++ b/packages/polywrap-msgpack/pyproject.toml @@ -48,6 +48,7 @@ disable = [ "too-many-return-statements", # too picky about return statements "protected-access", # Needed for internal use "invalid-name", # too picky about names + "too-many-branches", # sanitize function has too many branches ] ignore = [ "tests/" diff --git a/packages/polywrap-msgpack/tests/strategies/__init__.py b/packages/polywrap-msgpack/tests/strategies/__init__.py index 60f9667b..bdbd479b 100644 --- a/packages/polywrap-msgpack/tests/strategies/__init__.py +++ b/packages/polywrap-msgpack/tests/strategies/__init__.py @@ -1,2 +1,4 @@ from .basic_strategies import * -from .class_strategies import * \ No newline at end of file +from .class_strategies import * +from .enum_strategies import * +from .generic_map_strategies import * diff --git a/packages/polywrap-msgpack/tests/strategies/enum_strategies.py b/packages/polywrap-msgpack/tests/strategies/enum_strategies.py new file mode 100644 index 00000000..7b65fe4d --- /dev/null +++ b/packages/polywrap-msgpack/tests/strategies/enum_strategies.py @@ -0,0 +1,14 @@ +from enum import IntEnum +from hypothesis import strategies as st + + +class TestEnum(IntEnum): + Test0 = 0 + Test1 = 1 + Test2 = 2 + Test3 = 3 + + +def enum_st() -> st.SearchStrategy[IntEnum]: + """Define a strategy for generating valid `Enum`.""" + return st.sampled_from(list(TestEnum)) diff --git a/packages/polywrap-msgpack/tests/test_mirror.py b/packages/polywrap-msgpack/tests/test_mirror.py index b06adbc2..d2565f4e 100644 --- a/packages/polywrap-msgpack/tests/test_mirror.py +++ b/packages/polywrap-msgpack/tests/test_mirror.py @@ -24,12 +24,19 @@ valid_generic_map_st, ) +from .strategies.enum_strategies import enum_st + @given(scalar_st()) def test_mirror_scalar(s: Any): assert msgpack_decode(msgpack_encode(s)) == s +@given(enum_st()) +def test_mirror_enum(s: Any): + assert msgpack_decode(msgpack_encode(s)) == s + + @given(sequence_of_scalar_st()) def test_mirror_any_sequence_of_scalars(s: Sequence[Any]): assert msgpack_decode(msgpack_encode(s)) == list(s)