From 9ef5503728f84f32ed407ebb2cafa72cb070c877 Mon Sep 17 00:00:00 2001
From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>
Date: Sun, 23 Aug 2020 07:26:15 -0500
Subject: [PATCH] Small improvements to models.py

---
 src/betterproto/plugin/models.py | 33 ++++++++++++++++----------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/src/betterproto/plugin/models.py b/src/betterproto/plugin/models.py
index c4906de..65c50a1 100644
--- a/src/betterproto/plugin/models.py
+++ b/src/betterproto/plugin/models.py
@@ -153,6 +153,7 @@ class ProtoContentBase:
 
     path: List[int]
     comment_indent: int = 4
+    parent: Union["Messsage", "OutputTemplate"]
 
     def __post_init__(self):
         """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."""
         name = f"{self.py_name}"
         annotations = f": {self.annotation}"
+        field_args = ", ".join(
+            ([""] + self.betterproto_field_args) if self.betterproto_field_args else []
+        )
         betterproto_field_type = (
             f"betterproto.{self.field_type}_field({self.proto_obj.number}"
-            + f"{self.betterproto_field_args}"
+            + field_args
             + ")"
         )
         return name + annotations + " = " + betterproto_field_type
 
     @property
-    def betterproto_field_args(self):
-        args = ""
+    def betterproto_field_args(self) -> List[str]:
+        args = []
         if self.field_wraps:
-            args = args + f", wraps={self.field_wraps}"
+            args.append(f"wraps={self.field_wraps}")
         return args
 
     @property
@@ -473,10 +477,10 @@ class FieldCompiler(MessageCompiler):
 @dataclass
 class OneOfFieldCompiler(FieldCompiler):
     @property
-    def betterproto_field_args(self) -> "str":
+    def betterproto_field_args(self) -> List[str]:
         args = super().betterproto_field_args
         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
 
 
@@ -505,16 +509,13 @@ class MapEntryCompiler(FieldCompiler):
                     self.proto_v_type = self.proto_obj.Type.Name(nested.field[1].type)
         super().__post_init__()  # call FieldCompiler-> MessageCompiler __post_init__
 
-    def get_field_string(self, indent: int = 4) -> str:
-        """Construct string representation of this field."""
-        name = f"{self.py_name}"
-        annotations = f": {self.annotation}"
-        betterproto_field_type = (
-            f"betterproto.map_field("
-            f"{self.proto_obj.number}, betterproto.{self.proto_k_type}, "
-            f"betterproto.{self.proto_v_type})"
-        )
-        return name + annotations + " = " + betterproto_field_type
+    @property
+    def betterproto_field_args(self) -> List[str]:
+        return [f"betterproto.{self.proto_k_type}", f"betterproto.{self.proto_v_type}"]
+
+    @property
+    def field_type(self) -> str:
+        return "map"
 
     @property
     def annotation(self):