From 4f48950346d60b30d6dbda58b9cabf1c591cc9ed Mon Sep 17 00:00:00 2001 From: Yannick Dylla Date: Sun, 3 Oct 2021 17:15:32 +0200 Subject: [PATCH 1/3] map enum int's into python enums (#157) --- src/betterproto/__init__.py | 8 ++++++-- tests/inputs/enum/test_enum.py | 20 ++++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/src/betterproto/__init__.py b/src/betterproto/__init__.py index 05b8b7cd1..ef24ab955 100644 --- a/src/betterproto/__init__.py +++ b/src/betterproto/__init__.py @@ -835,8 +835,8 @@ def _get_field_default_gen(cls, field: dataclasses.Field) -> Any: else: return t elif issubclass(t, Enum): - # Enums always default to zero. - return int + cls = cls._cls_for(field) + return lambda: cls(0) # Enums always default to zero. elif t is datetime: # Offsets are relative to 1970-01-01T00:00:00Z return datetime_default_gen @@ -861,6 +861,10 @@ def _postprocess_single( elif meta.proto_type == TYPE_BOOL: # Booleans use a varint encoding, so convert it to true/false. value = value > 0 + elif meta.proto_type == TYPE_ENUM: + # Convert enum ints to python enum instances + cls = self._betterproto.cls_by_field[field_name] + value = cls(value) elif wire_type in (WIRE_FIXED_32, WIRE_FIXED_64): fmt = _pack_fmt(meta.proto_type) value = struct.unpack(fmt, value)[0] diff --git a/tests/inputs/enum/test_enum.py b/tests/inputs/enum/test_enum.py index 3005c43a8..bc3c8c961 100644 --- a/tests/inputs/enum/test_enum.py +++ b/tests/inputs/enum/test_enum.py @@ -82,3 +82,23 @@ def enum_generator(): yield Choice.THREE assert Test(choices=enum_generator()).to_dict()["choices"] == ["ONE", "THREE"] + + +def test_enum_mapped_on_parse(): + # test default value + b = Test().parse(bytes(Test())) + assert b.choice.name == Choice.ZERO.name + assert b.choices == [] + + # test non default value + a = Test().parse(bytes(Test(choice=Choice.ONE))) + assert a.choice.name == Choice.ONE.name + assert b.choices == [] + + # test repeated + c = Test().parse(bytes(Test(choices=[Choice.THREE, Choice.FOUR]))) + assert c.choices[0].name == Choice.THREE.name + assert c.choices[1].name == Choice.FOUR.name + + # bonus: defaults after empty init are also mapped + assert Test().choice.name == Choice.ZERO.name From 392d00990143246be00059e81898f82eea0442fc Mon Sep 17 00:00:00 2001 From: Yannick Dylla Date: Sun, 3 Oct 2021 18:52:25 +0200 Subject: [PATCH 2/3] fix enum not having value for zero Fallback to int like before, not really happy with this but no other idea. --- src/betterproto/__init__.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/betterproto/__init__.py b/src/betterproto/__init__.py index ef24ab955..0de04652b 100644 --- a/src/betterproto/__init__.py +++ b/src/betterproto/__init__.py @@ -835,8 +835,15 @@ def _get_field_default_gen(cls, field: dataclasses.Field) -> Any: else: return t elif issubclass(t, Enum): - cls = cls._cls_for(field) - return lambda: cls(0) # Enums always default to zero. + # Enums always default to zero. + def default_enum(): + try: + # try to create a python enum instance + return t(0) + except ValueError: + return 0 # if that does not work fallback to int + + return default_enum elif t is datetime: # Offsets are relative to 1970-01-01T00:00:00Z return datetime_default_gen From 721c7256b7b296804e21f4cb483371872ec7b62b Mon Sep 17 00:00:00 2001 From: Yannick Dylla Date: Mon, 18 Oct 2021 15:49:29 +0200 Subject: [PATCH 3/3] fallback to int when received value does not exists in enum --- src/betterproto/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/betterproto/__init__.py b/src/betterproto/__init__.py index 0de04652b..024ea19fd 100644 --- a/src/betterproto/__init__.py +++ b/src/betterproto/__init__.py @@ -871,7 +871,10 @@ def _postprocess_single( elif meta.proto_type == TYPE_ENUM: # Convert enum ints to python enum instances cls = self._betterproto.cls_by_field[field_name] - value = cls(value) + try: + value = cls(value) + except ValueError: + pass # the received value does not exist in the enum so we have to pass it as raw int elif wire_type in (WIRE_FIXED_32, WIRE_FIXED_64): fmt = _pack_fmt(meta.proto_type) value = struct.unpack(fmt, value)[0]