Move parsing of protobuf data types and services into separate methods

betterproto/plugin.py

@@ -163,13 +163,104 @@ def generate_code(request, response):
         for proto_file in output_package_content["files"]:
             item: DescriptorProto
             for item, path in traverse(proto_file):
-                data = {"name": item.name, "py_name": pythonize_class_name(item.name)}
+                read_protobuf_type(input_package_name, item, path, proto_file, template_data)
 
+            for i, service in enumerate(proto_file.service):
+                read_protobuf_service(i, input_package_name, proto_file, service, template_data)
+
+        template_data["imports"] = sorted(template_data["imports"])
+        template_data["datetime_imports"] = sorted(template_data["datetime_imports"])
+        template_data["typing_imports"] = sorted(template_data["typing_imports"])
+
+        # Fill response
+        output_path = pathlib.Path(*output_package_name.split("."), "__init__.py")
+        output_paths.add(output_path)
+
+        f = response.file.add()
+        f.name = str(output_path)
+
+        # Render and then format the output file.
+        f.content = black.format_str(
+            template.render(description=template_data),
+            mode=black.FileMode(target_versions=set([black.TargetVersion.PY37])),
+        )
+
+    # Make each output directory a package with __init__ file
+    init_files = (
+        set(
+            directory.joinpath("__init__.py")
+            for path in output_paths
+            for directory in path.parents
+        )
+        - output_paths
+    )
+
+    for init_file in init_files:
+        init = response.file.add()
+        init.name = str(init_file)
+
+    for output_package_name in sorted(output_paths.union(init_files)):
+        print(f"Writing {output_package_name}", file=sys.stderr)
+
+
+def read_protobuf_service(i, input_package_name, proto_file, service, template_data):
+    # print(service, file=sys.stderr)
+    data = {
+        "name": service.name,
+        "py_name": pythonize_class_name(service.name),
+        "comment": get_comment(proto_file, [6, i]),
+        "methods": [],
+    }
+    for j, method in enumerate(service.method):
+        input_message = None
+        input_type = get_type_reference(
+            input_package_name, template_data["imports"], method.input_type
+        ).strip('"')
+        for msg in template_data["messages"]:
+            if msg["name"] == input_type:
+                input_message = msg
+                for field in msg["properties"]:
+                    if field["zero"] == "None":
+                        template_data["typing_imports"].add("Optional")
+                break
+
+        data["methods"].append(
+            {
+                "name": method.name,
+                "py_name": pythonize_method_name(method.name),
+                "comment": get_comment(proto_file, [6, i, 2, j], indent=8),
+                "route": f"/{input_package_name}.{service.name}/{method.name}",
+                "input": get_type_reference(
+                    input_package_name, template_data["imports"], method.input_type
+                ).strip('"'),
+                "input_message": input_message,
+                "output": get_type_reference(
+                    input_package_name,
+                    template_data["imports"],
+                    method.output_type,
+                    unwrap=False,
+                ),
+                "client_streaming": method.client_streaming,
+                "server_streaming": method.server_streaming,
+            }
+        )
+
+        if method.client_streaming:
+            template_data["typing_imports"].add("AsyncIterable")
+            template_data["typing_imports"].add("Iterable")
+            template_data["typing_imports"].add("Union")
+        if method.server_streaming:
+            template_data["typing_imports"].add("AsyncIterator")
+    template_data["services"].append(data)
+
+
+def read_protobuf_type(input_package_name, item, path, proto_file, template_data):
+    data = {"name": item.name, "py_name": pythonize_class_name(item.name)}
     if isinstance(item, DescriptorProto):
         # print(item, file=sys.stderr)
         if item.options.map_entry:
             # Skip generated map entry messages since we just use dicts
-                        continue
+            return
 
         data.update(
             {
@@ -291,93 +382,6 @@ def generate_code(request, response):
 
         template_data["enums"].append(data)
 
-            for i, service in enumerate(proto_file.service):
-                # print(service, file=sys.stderr)
-
-                data = {
-                    "name": service.name,
-                    "py_name": pythonize_class_name(service.name),
-                    "comment": get_comment(proto_file, [6, i]),
-                    "methods": [],
-                }
-
-                for j, method in enumerate(service.method):
-                    input_message = None
-                    input_type = get_type_reference(
-                        input_package_name, template_data["imports"], method.input_type
-                    ).strip('"')
-                    for msg in template_data["messages"]:
-                        if msg["name"] == input_type:
-                            input_message = msg
-                            for field in msg["properties"]:
-                                if field["zero"] == "None":
-                                    template_data["typing_imports"].add("Optional")
-                            break
-
-                    data["methods"].append(
-                        {
-                            "name": method.name,
-                            "py_name": pythonize_method_name(method.name),
-                            "comment": get_comment(proto_file, [6, i, 2, j], indent=8),
-                            "route": f"/{input_package_name}.{service.name}/{method.name}",
-                            "input": get_type_reference(
-                                input_package_name, template_data["imports"], method.input_type
-                            ).strip('"'),
-                            "input_message": input_message,
-                            "output": get_type_reference(
-                                input_package_name,
-                                template_data["imports"],
-                                method.output_type,
-                                unwrap=False,
-                            ),
-                            "client_streaming": method.client_streaming,
-                            "server_streaming": method.server_streaming,
-                        }
-                    )
-
-                    if method.client_streaming:
-                        template_data["typing_imports"].add("AsyncIterable")
-                        template_data["typing_imports"].add("Iterable")
-                        template_data["typing_imports"].add("Union")
-                    if method.server_streaming:
-                        template_data["typing_imports"].add("AsyncIterator")
-
-                template_data["services"].append(data)
-
-        template_data["imports"] = sorted(template_data["imports"])
-        template_data["datetime_imports"] = sorted(template_data["datetime_imports"])
-        template_data["typing_imports"] = sorted(template_data["typing_imports"])
-
-        # Fill response
-        output_path = pathlib.Path(*output_package_name.split("."), "__init__.py")
-        output_paths.add(output_path)
-
-        f = response.file.add()
-        f.name = str(output_path)
-
-        # Render and then format the output file.
-        f.content = black.format_str(
-            template.render(description=template_data),
-            mode=black.FileMode(target_versions=set([black.TargetVersion.PY37])),
-        )
-
-    # Make each output directory a package with __init__ file
-    init_files = (
-        set(
-            directory.joinpath("__init__.py")
-            for path in output_paths
-            for directory in path.parents
-        )
-        - output_paths
-    )
-
-    for init_file in init_files:
-        init = response.file.add()
-        init.name = str(init_file)
-
-    for output_package_name in sorted(output_paths.union(init_files)):
-        print(f"Writing {output_package_name}", file=sys.stderr)
-
 
 def main():
     """The plugin's main entry point."""