Implement imports, simplified default value handling

Pipfile

@@ -18,5 +18,6 @@ jinja2 = "*"
 python_version = "3.7"
 
 [scripts]
+plugin = "protoc --plugin=protoc-gen-custom=protoc-gen-betterpy.py --custom_out=."
 generate = "python betterproto/tests/generate.py"
 test = "pytest ./betterproto/tests"

README.md

@@ -13,6 +13,7 @@ This project aims to provide an improved experience when using Protobuf / gRPC i
 
 This project is heavily inspired by, and borrows functionality from:
 
+- https://github.com/protocolbuffers/protobuf/tree/master/python
 - https://github.com/eigenein/protobuf/
 - https://github.com/vmagamedov/grpclib
 
@@ -27,8 +28,8 @@ This project is heavily inspired by, and borrows functionality from:
 - [x] Maps
   - [x] Maps of message fields
 - [ ] Support passthrough of unknown fields
-- [ ] Refs to nested types
+- [x] Refs to nested types
-- [ ] Imports in proto files
+- [x] Imports in proto files
 - [ ] Well-known Google types
 - [ ] JSON that isn't completely naive.
 - [ ] Async service stubs

betterproto/__init__.py

@@ -92,6 +92,18 @@ WIRE_FIXED_64_TYPES = [TYPE_DOUBLE, TYPE_FIXED64, TYPE_SFIXED64]
 WIRE_LEN_DELIM_TYPES = [TYPE_STRING, TYPE_BYTES, TYPE_MESSAGE, TYPE_MAP]
 
 
+def get_default(proto_type: int) -> Any:
+    """Get the default (zero value) for a given type."""
+    return {
+        TYPE_BOOL: False,
+        TYPE_FLOAT: 0.0,
+        TYPE_DOUBLE: 0.0,
+        TYPE_STRING: "",
+        TYPE_BYTES: b"",
+        TYPE_MAP: {},
+    }.get(proto_type, 0)
+
+
 @dataclasses.dataclass(frozen=True)
 class FieldMetadata:
     """Stores internal metadata used for parsing & serialization."""
@@ -114,7 +126,7 @@ class FieldMetadata:
 def dataclass_field(
     number: int,
     proto_type: str,
-    default: Any,
+    default: Any = None,
     map_types: Optional[Tuple[str, str]] = None,
     **kwargs: dict,
 ) -> dataclasses.Field:
@@ -141,6 +153,10 @@ def enum_field(number: int, default: Union[int, Type[Iterable]] = 0) -> Any:
     return dataclass_field(number, TYPE_ENUM, default=default)
 
 
+def bool_field(number: int, default: Union[bool, Type[Iterable]] = 0) -> Any:
+    return dataclass_field(number, TYPE_BOOL, default=default)
+
+
 def int32_field(number: int, default: Union[int, Type[Iterable]] = 0) -> Any:
     return dataclass_field(number, TYPE_INT32, default=default)
 
@@ -193,8 +209,8 @@ def string_field(number: int, default: str = "") -> Any:
     return dataclass_field(number, TYPE_STRING, default=default)
 
 
-def message_field(number: int, default: Type["Message"]) -> Any:
+def message_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_MESSAGE, default=default)
+    return dataclass_field(number, TYPE_MESSAGE)
 
 
 def map_field(number: int, key_type: str, value_type: str) -> Any:
@@ -345,6 +361,29 @@ class Message(ABC):
     to go between Python, binary and JSON protobuf message representations.
     """
 
+    def __post_init__(self) -> None:
+        # Set a default value for each field in the class after `__init__` has
+        # already been run.
+        for field in dataclasses.fields(self):
+            meta = FieldMetadata.get(field)
+
+            t = self._cls_for(field, index=-1)
+
+            value = 0
+            if meta.proto_type == TYPE_MAP:
+                # Maps cannot be repeated, so we check these first.
+                value = {}
+            elif hasattr(t, "__args__") and len(t.__args__) == 1:
+                # Anything else with type args is a list.
+                value = []
+            elif meta.proto_type == TYPE_MESSAGE:
+                # Message means creating an instance of the right type.
+                value = t()
+            else:
+                value = get_default(meta.proto_type)
+
+            setattr(self, field.name, value)
+
     def __bytes__(self) -> bytes:
         """
         Get the binary encoded Protobuf representation of this instance.
@@ -356,6 +395,7 @@ class Message(ABC):
 
             if isinstance(value, list):
                 if not len(value):
+                    # Empty values are not serialized
                     continue
 
                 if meta.proto_type in PACKED_TYPES:
@@ -371,6 +411,7 @@ class Message(ABC):
                         output += _serialize_single(meta.number, meta.proto_type, item)
             elif isinstance(value, dict):
                 if not len(value):
+                    # Empty values are not serialized
                     continue
 
                 for k, v in value.items():
@@ -378,7 +419,8 @@ class Message(ABC):
                     sv = _serialize_single(2, meta.map_types[1], v)
                     output += _serialize_single(meta.number, meta.proto_type, sk + sv)
             else:
-                if value == field.default:
+                if value == get_default(meta.proto_type):
+                    # Default (zero) values are not serialized
                     continue
 
                 output += _serialize_single(meta.number, meta.proto_type, value)
@@ -390,7 +432,7 @@ class Message(ABC):
         module = inspect.getmodule(self)
         type_hints = get_type_hints(self, vars(module))
         cls = type_hints[field.name]
-        if hasattr(cls, "__args__"):
+        if hasattr(cls, "__args__") and index >= 0:
             cls = type_hints[field.name].__args__[index]
         return cls
 
@@ -522,7 +564,7 @@ class Message(ABC):
         """
         for field in dataclasses.fields(self):
             meta = FieldMetadata.get(field)
-            if field.name in value:
+            if field.name in value and value[field.name] is not None:
                 if meta.proto_type == "message":
                     v = getattr(self, field.name)
                     # print(v, value[field.name])

betterproto/templates/main.py

@@ -7,6 +7,10 @@ from dataclasses import dataclass
 from typing import Dict, List
 
 import betterproto
+{% for i in description.imports %}
+
+{{ i }}
+{% endfor %}
 
 
 {% if description.enums %}{% for enum in description.enums %}
@@ -21,9 +25,9 @@ class {{ enum.name }}(enum.IntEnum):
         {% endif %}
     {{ entry.name }} = {{ entry.value }}
     {% endfor %}
+
+
 {% endfor %}
-
-
 {% endif %}
 {% for message in description.messages %}
 @dataclass
@@ -36,8 +40,11 @@ class {{ message.name }}(betterproto.Message):
         {% if field.comment %}
 {{ field.comment }}
         {% endif %}
-    {{ field.name }}: {{ field.type }} = betterproto.{{ field.field_type }}_field({{ field.number }}{% if field.zero and field.field_type != 'map' %}, default={{ field.zero }}{% endif %}{% if field.field_type == 'map'%}, betterproto.{{ field.map_types[0] }}, betterproto.{{ field.map_types[1] }}{% endif %})
+    {{ field.name }}: {{ field.type }} = betterproto.{{ field.field_type }}_field({{ field.number }}{% if field.field_type == 'map'%}, betterproto.{{ field.map_types[0] }}, betterproto.{{ field.map_types[1] }}{% endif %})
     {% endfor %}
+    {% if not message.properties %}
+    pass
+    {% endif %}
 
 
 {% endfor %}

betterproto/tests/ref.json

@@ -0,0 +1,5 @@
+{
+  "greeting": {
+    "greeting": "hello"
+  }
+}

betterproto/tests/ref.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+package ref;
+
+import "repeatedmessage.proto";
+
+message Test {
+  repeatedmessage.Sub greeting = 1;
+}

betterproto/tests/repeatedmessage.proto

@@ -1,5 +1,7 @@
 syntax = "proto3";
 
+package repeatedmessage;
+
 message Test {
   repeated Sub greetings = 1;
 }

betterproto/tests/test_inputs.py

@@ -2,7 +2,7 @@ import importlib
 import pytest
 import json
 
-from generate import get_files, get_base
+from .generate import get_files, get_base
 
 inputs = get_files(".bin")
 
@@ -10,7 +10,7 @@ inputs = get_files(".bin")
 @pytest.mark.parametrize("filename", inputs)
 def test_sample(filename: str) -> None:
     module = get_base(filename).split("-")[0]
-    imported = importlib.import_module(module)
+    imported = importlib.import_module(f"betterproto.tests.{module}")
     data_binary = open(filename, "rb").read()
     data_dict = json.loads(open(filename.replace(".bin", ".json")).read())
     t1 = imported.Test().parse(data_binary)

protoc-gen-betterpy.py

@@ -22,33 +22,45 @@ from jinja2 import Environment, PackageLoader
 
 
 def py_type(
-    message: DescriptorProto, descriptor: FieldDescriptorProto
+    package: str,
-) -> Tuple[str, str]:
+    imports: set,
+    message: DescriptorProto,
+    descriptor: FieldDescriptorProto,
+) -> str:
     if descriptor.type in [1, 2, 6, 7, 15, 16]:
-        return "float", descriptor.default_value
+        return "float"
     elif descriptor.type in [3, 4, 5, 13, 17, 18]:
-        return "int", descriptor.default_value
+        return "int"
     elif descriptor.type == 8:
-        return "bool", descriptor.default_value.capitalize()
+        return "bool"
     elif descriptor.type == 9:
-        default = ""
+        return "str"
-        if descriptor.default_value:
+    elif descriptor.type in [11, 14]:
-            default = f'"{descriptor.default_value}"'
+        # Type referencing another defined Message or a named enum
-        return "str", default
+        message_type = descriptor.type_name.lstrip(".")
-    elif descriptor.type == 11:
+        if message_type.startswith(package):
-        # Type referencing another defined Message
+            # This is the current package, which has nested types flattened.
-        # print(descriptor.type_name, file=sys.stderr)
+            message_type = message_type.lstrip(package).lstrip(".").replace(".", "")
-        # message_type = descriptor.type_name.replace(".", "")
+
-        message_type = descriptor.type_name.split(".").pop()
+        if "." in message_type:
-        return f'"{message_type}"', f"lambda: {message_type}()"
+            # This is imported from another package. No need
+            # to use a forward ref and we need to add the import.
+            message_type = message_type.strip('"')
+            parts = message_type.split(".")
+            imports.add(f"from .{'.'.join(parts[:-2])} import {parts[-2]}")
+            message_type = f"{parts[-2]}.{parts[-1]}"
+
+        # print(
+        #     descriptor.name,
+        #     package,
+        #     descriptor.type_name,
+        #     message_type,
+        #     file=sys.stderr,
+        # )
+
+        return f'"{message_type}"'
     elif descriptor.type == 12:
-        default = ""
+        return "bytes"
-        if descriptor.default_value:
-            default = f'b"{descriptor.default_value}"'
-        return "bytes", default
-    elif descriptor.type == 14:
-        # print(descriptor.type_name, file=sys.stderr)
-        return descriptor.type_name.split(".").pop(), 0
     else:
         raise NotImplementedError(f"Unknown type {descriptor.type}")
 
@@ -64,6 +76,8 @@ def traverse(proto_file):
 
                 if item.nested_type:
                     for n, p in _traverse(path + [i, 3], item.nested_type):
+                        # Adjust the name since we flatten the heirarchy.
+                        n.name = item.name + n.name
                         yield n, p
 
     return itertools.chain(
@@ -85,6 +99,7 @@ def get_comment(proto_file, path: List[int]) -> str:
             else:
                 # This is a class
                 if len(lines) == 1 and len(lines[0]) < 70:
+                    lines[0] = lines[0].strip('"')
                     return f'    """{lines[0]}"""'
                 else:
                     return f'    """\n{"    ".join(lines)}\n    """'
@@ -100,17 +115,33 @@ def generate_code(request, response):
     )
     template = env.get_template("main.py")
 
+    # TODO: Refactor below to generate a single file per package if packages
+    # are being used, otherwise one output for each input. Figure out how to
+    # set up relative imports when needed and change the Message type refs to
+    # use the import names when not in the current module.
+    output_map = {}
     for proto_file in request.proto_file:
-        # print(proto_file.message_type, file=sys.stderr)
+        out = proto_file.package
-        # print(proto_file.source_code_info, file=sys.stderr)
+        if not out:
-        output = {
+            out = os.path.splitext(proto_file.name)[0].replace(os.path.sep, ".")
-            "package": proto_file.package,
+
-            "filename": proto_file.name,
+        if out not in output_map:
-            "messages": [],
+            output_map[out] = {"package": proto_file.package, "files": []}
-            "enums": [],
+        output_map[out]["files"].append(proto_file)
-        }
+
+    # TODO: Figure out how to handle gRPC request/response messages and add
+    # processing below for Service.
+
+    for filename, options in output_map.items():
+        package = options["package"]
+        # print(package, filename, file=sys.stderr)
+        output = {"package": package, "imports": set(), "messages": [], "enums": []}
+
+        for proto_file in options["files"]:
+            # print(proto_file.message_type, file=sys.stderr)
+            # print(proto_file.service, file=sys.stderr)
+            # print(proto_file.source_code_info, file=sys.stderr)
 
-        # Parse request
             for item, path in traverse(proto_file):
                 # print(item, file=sys.stderr)
                 # print(path, file=sys.stderr)
@@ -131,7 +162,8 @@ def generate_code(request, response):
                     )
 
                     for i, f in enumerate(item.field):
-                    t, zero = py_type(item, f)
+                        t = py_type(package, output["imports"], item, f)
+
                         repeated = False
                         packed = False
 
@@ -146,11 +178,20 @@ def generate_code(request, response):
                                 for nested in item.nested_type:
                                     if nested.name == map_entry:
                                         if nested.options.map_entry:
-                                        print("Found a map!", file=sys.stderr)
+                                            # print("Found a map!", file=sys.stderr)
-                                        k, _ = py_type(item, nested.field[0])
+                                            k = py_type(
-                                        v, _ = py_type(item, nested.field[1])
+                                                package,
+                                                output["imports"],
+                                                item,
+                                                nested.field[0],
+                                            )
+                                            v = py_type(
+                                                package,
+                                                output["imports"],
+                                                item,
+                                                nested.field[1],
+                                            )
                                             t = f"Dict[{k}, {v}]"
-                                        zero = "dict"
                                             field_type = "map"
                                             map_types = (
                                                 f.Type.Name(nested.field[0].type),
@@ -161,7 +202,6 @@ def generate_code(request, response):
                             # Repeated field
                             repeated = True
                             t = f"List[{t}]"
-                        zero = "list"
 
                             if f.type in [1, 2, 3, 4, 5, 6, 7, 8, 13, 15, 16, 17, 18]:
                                 packed = True
@@ -175,7 +215,6 @@ def generate_code(request, response):
                                 "field_type": field_type,
                                 "map_types": map_types,
                                 "type": t,
-                            "zero": zero,
                                 "repeated": repeated,
                                 "packed": packed,
                             }
@@ -203,9 +242,12 @@ def generate_code(request, response):
 
                     output["enums"].append(data)
 
+        output["imports"] = sorted(output["imports"])
+
         # Fill response
         f = response.file.add()
-        f.name = os.path.splitext(proto_file.name)[0] + ".py"
+        # print(filename, file=sys.stderr)
+        f.name = filename + ".py"
         # f.content = json.dumps(output, indent=2)
         f.content = template.render(description=output).rstrip("\n") + "\n"