Skip to content

Commit 42a1d7f

Browse files
adriangbGobot1234
authored andcommitted
Small improvements to models.py
1 parent bf0ca78 commit 42a1d7f

File tree

1 file changed

+17
-16
lines changed

1 file changed

+17
-16
lines changed

src/betterproto/plugin/models.py

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,7 @@ class ProtoContentBase:
153153

154154
path: List[int]
155155
comment_indent: int = 4
156+
parent: Union["Messsage", "OutputTemplate"]
156157

157158
def __post_init__(self):
158159
"""Checks that no fake default fields were left as placeholders."""
@@ -355,18 +356,21 @@ def get_field_string(self, indent: int = 4) -> str:
355356
"""Construct string representation of this field as a field."""
356357
name = f"{self.py_name}"
357358
annotations = f": {self.annotation}"
359+
field_args = ", ".join(
360+
([""] + self.betterproto_field_args) if self.betterproto_field_args else []
361+
)
358362
betterproto_field_type = (
359363
f"betterproto.{self.field_type}_field({self.proto_obj.number}"
360-
+ f"{self.betterproto_field_args}"
364+
+ field_args
361365
+ ")"
362366
)
363367
return name + annotations + " = " + betterproto_field_type
364368

365369
@property
366-
def betterproto_field_args(self):
367-
args = ""
370+
def betterproto_field_args(self) -> List[str]:
371+
args = []
368372
if self.field_wraps:
369-
args = args + f", wraps={self.field_wraps}"
373+
args.append(f"wraps={self.field_wraps}")
370374
return args
371375

372376
@property
@@ -473,10 +477,10 @@ def annotation(self) -> str:
473477
@dataclass
474478
class OneOfFieldCompiler(FieldCompiler):
475479
@property
476-
def betterproto_field_args(self) -> "str":
480+
def betterproto_field_args(self) -> List[str]:
477481
args = super().betterproto_field_args
478482
group = self.parent.proto_obj.oneof_decl[self.proto_obj.oneof_index].name
479-
args = args + f', group="{group}"'
483+
args.append(f'group="{group}"')
480484
return args
481485

482486

@@ -505,16 +509,13 @@ def __post_init__(self):
505509
self.proto_v_type = self.proto_obj.Type.Name(nested.field[1].type)
506510
super().__post_init__() # call FieldCompiler-> MessageCompiler __post_init__
507511

508-
def get_field_string(self, indent: int = 4) -> str:
509-
"""Construct string representation of this field."""
510-
name = f"{self.py_name}"
511-
annotations = f": {self.annotation}"
512-
betterproto_field_type = (
513-
f"betterproto.map_field("
514-
f"{self.proto_obj.number}, betterproto.{self.proto_k_type}, "
515-
f"betterproto.{self.proto_v_type})"
516-
)
517-
return name + annotations + " = " + betterproto_field_type
512+
@property
513+
def betterproto_field_args(self) -> List[str]:
514+
return [f"betterproto.{self.proto_k_type}", f"betterproto.{self.proto_v_type}"]
515+
516+
@property
517+
def field_type(self) -> str:
518+
return "map"
518519

519520
@property
520521
def annotation(self):

0 commit comments

Comments
 (0)