Fix default values for enum service args #298

Author: Kim Gustyr

src/betterproto/plugin/models.py

@@ -59,8 +59,7 @@
 import re
 import textwrap
 from dataclasses import dataclass, field
-from typing import Dict, Iterable, Iterator, List, Optional, Set, Text, Type, Union
-import sys
+from typing import Dict, Iterable, Iterator, List, Optional, Set, Type, Union
 
 from ..casing import sanitize_name
 from ..compile.importing import get_type_reference, parse_source_type_name
@@ -460,7 +459,7 @@ def field_type(self) -> str:
         )
 
     @property
-    def default_value_string(self) -> Union[Text, None, float, int]:
+    def default_value_string(self) -> str:
         """Python representation of the default proto value."""
         if self.repeated:
             return "[]"
@@ -474,6 +473,10 @@ def default_value_string(self) -> Union[Text, None, float, int]:
             return '""'
         elif self.py_type == "bytes":
             return 'b""'
+        elif self.field_type == "enum":
+            enum_proto_obj_name = self.proto_obj.type_name.split(".").pop()
+            enum = next(e for e in self.output_file.enums if e.proto_obj.name == enum_proto_obj_name)
+            return enum.default_value_string
         else:
             # Message type
             return "None"

tests/inputs/service/service.proto

@@ -2,9 +2,16 @@ syntax = "proto3";
 
 package service;
 
+enum ThingType {
+  UNKNOWN = 0;
+  LIVING = 1;
+  DEAD = 2;
+}
+
 message DoThingRequest {
   string name = 1;
   repeated string comments = 2;
+  ThingType type = 3;
 }
 
 message DoThingResponse {

tests/test_features.py

@@ -1,7 +1,8 @@
 import betterproto
 from dataclasses import dataclass
 from typing import Optional, List, Dict
-from datetime import datetime, timedelta
+from datetime import datetime
+from inspect import signature
 
 
 def test_has_field():
@@ -476,3 +477,10 @@ class Envelope(betterproto.Message):
 
     msg.from_dict({"timestamps": iso_candidates})
     assert all([isinstance(item, datetime) for item in msg.timestamps])
+
+
+def test_enum_service_argument__expected_default_value():
+    from tests.output_betterproto.service.service import ThingType, TestStub
+
+    sig = signature(TestStub.do_thing)
+    assert sig.parameters["type"].default == ThingType.UNKNOWN