map enum int's into python enums (danielgtaylor#157)

Author: ydylla

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]

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