Fix default values for enum service args #298 (#299)

This commit is contained in:
Kim Gustyr 2021-12-04 00:26:48 +03:00 committed by GitHub
parent b0a36d12e4
commit bc3cfc5562
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 26 additions and 4 deletions

View File

@ -59,8 +59,7 @@ from betterproto.lib.google.protobuf.compiler import CodeGeneratorRequest
import re import re
import textwrap import textwrap
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Dict, Iterable, Iterator, List, Optional, Set, Text, Type, Union from typing import Dict, Iterable, Iterator, List, Optional, Set, Type, Union
import sys
from ..casing import sanitize_name from ..casing import sanitize_name
from ..compile.importing import get_type_reference, parse_source_type_name from ..compile.importing import get_type_reference, parse_source_type_name
@ -460,7 +459,7 @@ class FieldCompiler(MessageCompiler):
) )
@property @property
def default_value_string(self) -> Union[Text, None, float, int]: def default_value_string(self) -> str:
"""Python representation of the default proto value.""" """Python representation of the default proto value."""
if self.repeated: if self.repeated:
return "[]" return "[]"
@ -474,6 +473,14 @@ class FieldCompiler(MessageCompiler):
return '""' return '""'
elif self.py_type == "bytes": elif self.py_type == "bytes":
return 'b""' 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: else:
# Message type # Message type
return "None" return "None"

View File

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

View File

@ -1,7 +1,8 @@
import betterproto import betterproto
from dataclasses import dataclass from dataclasses import dataclass
from typing import Optional, List, Dict from typing import Optional, List, Dict
from datetime import datetime, timedelta from datetime import datetime
from inspect import signature
def test_has_field(): def test_has_field():
@ -476,3 +477,10 @@ def test_iso_datetime_list():
msg.from_dict({"timestamps": iso_candidates}) msg.from_dict({"timestamps": iso_candidates})
assert all([isinstance(item, datetime) for item in msg.timestamps]) 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