Small improvements to models.py

This commit is contained in:
Adrian Garcia Badaracco 2020-08-23 07:26:15 -05:00 committed by GitHub
parent c93351ef21
commit 9ef5503728
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -153,6 +153,7 @@ class ProtoContentBase:
path: List[int] path: List[int]
comment_indent: int = 4 comment_indent: int = 4
parent: Union["Messsage", "OutputTemplate"]
def __post_init__(self): def __post_init__(self):
"""Checks that no fake default fields were left as placeholders.""" """Checks that no fake default fields were left as placeholders."""
@ -355,18 +356,21 @@ class FieldCompiler(MessageCompiler):
"""Construct string representation of this field as a field.""" """Construct string representation of this field as a field."""
name = f"{self.py_name}" name = f"{self.py_name}"
annotations = f": {self.annotation}" annotations = f": {self.annotation}"
field_args = ", ".join(
([""] + self.betterproto_field_args) if self.betterproto_field_args else []
)
betterproto_field_type = ( betterproto_field_type = (
f"betterproto.{self.field_type}_field({self.proto_obj.number}" f"betterproto.{self.field_type}_field({self.proto_obj.number}"
+ f"{self.betterproto_field_args}" + field_args
+ ")" + ")"
) )
return name + annotations + " = " + betterproto_field_type return name + annotations + " = " + betterproto_field_type
@property @property
def betterproto_field_args(self): def betterproto_field_args(self) -> List[str]:
args = "" args = []
if self.field_wraps: if self.field_wraps:
args = args + f", wraps={self.field_wraps}" args.append(f"wraps={self.field_wraps}")
return args return args
@property @property
@ -473,10 +477,10 @@ class FieldCompiler(MessageCompiler):
@dataclass @dataclass
class OneOfFieldCompiler(FieldCompiler): class OneOfFieldCompiler(FieldCompiler):
@property @property
def betterproto_field_args(self) -> "str": def betterproto_field_args(self) -> List[str]:
args = super().betterproto_field_args args = super().betterproto_field_args
group = self.parent.proto_obj.oneof_decl[self.proto_obj.oneof_index].name group = self.parent.proto_obj.oneof_decl[self.proto_obj.oneof_index].name
args = args + f', group="{group}"' args.append(f'group="{group}"')
return args return args
@ -505,16 +509,13 @@ class MapEntryCompiler(FieldCompiler):
self.proto_v_type = self.proto_obj.Type.Name(nested.field[1].type) self.proto_v_type = self.proto_obj.Type.Name(nested.field[1].type)
super().__post_init__() # call FieldCompiler-> MessageCompiler __post_init__ super().__post_init__() # call FieldCompiler-> MessageCompiler __post_init__
def get_field_string(self, indent: int = 4) -> str: @property
"""Construct string representation of this field.""" def betterproto_field_args(self) -> List[str]:
name = f"{self.py_name}" return [f"betterproto.{self.proto_k_type}", f"betterproto.{self.proto_v_type}"]
annotations = f": {self.annotation}"
betterproto_field_type = ( @property
f"betterproto.map_field(" def field_type(self) -> str:
f"{self.proto_obj.number}, betterproto.{self.proto_k_type}, " return "map"
f"betterproto.{self.proto_v_type})"
)
return name + annotations + " = " + betterproto_field_type
@property @property
def annotation(self): def annotation(self):