From 315a823c865bd47858a37fa91ec4669d80d51e70 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Michal=20August=C3=BDn?= Date: Tue, 25 Oct 2022 10:14:22 +0200 Subject: [PATCH] feat: int mapped to enum --- src/betterproto/__init__.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/betterproto/__init__.py b/src/betterproto/__init__.py index e6af1bcb8..4bea15a86 100644 --- a/src/betterproto/__init__.py +++ b/src/betterproto/__init__.py @@ -904,7 +904,14 @@ def _get_field_default_gen(cls, field: dataclasses.Field) -> Any: return t elif issubclass(t, Enum): # Enums always default to zero. - return int + 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 @@ -929,6 +936,13 @@ 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] + 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]