Add OneOf support, rework field detection

betterproto/__init__.py

@@ -130,6 +130,8 @@ class FieldMetadata:
     proto_type: str
     # Map information if the proto_type is a map
     map_types: Optional[Tuple[str, str]]
+    # Groups several "one-of" fields together
+    group: Optional[str]
 
     @staticmethod
     def get(field: dataclasses.Field) -> "FieldMetadata":
@@ -138,12 +140,16 @@ class FieldMetadata:
 
 
 def dataclass_field(
-    number: int, proto_type: str, map_types: Optional[Tuple[str, str]] = None
+    number: int,
+    proto_type: str,
+    *,
+    map_types: Optional[Tuple[str, str]] = None,
+    group: Optional[str] = None,
 ) -> dataclasses.Field:
     """Creates a dataclass field with attached protobuf metadata."""
     return dataclasses.field(
         default=PLACEHOLDER,
-        metadata={"betterproto": FieldMetadata(number, proto_type, map_types)},
+        metadata={"betterproto": FieldMetadata(number, proto_type, map_types, group)},
     )
 
 
@@ -152,76 +158,80 @@ def dataclass_field(
 # out at runtime. The generated dataclass variables are still typed correctly.
 
 
-def enum_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_ENUM)
+def enum_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_ENUM, group=group)
 
 
-def bool_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_BOOL)
+def bool_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_BOOL, group=group)
 
 
-def int32_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_INT32)
+def int32_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_INT32, group=group)
 
 
-def int64_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_INT64)
+def int64_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_INT64, group=group)
 
 
-def uint32_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_UINT32)
+def uint32_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_UINT32, group=group)
 
 
-def uint64_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_UINT64)
+def uint64_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_UINT64, group=group)
 
 
-def sint32_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_SINT32)
+def sint32_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_SINT32, group=group)
 
 
-def sint64_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_SINT64)
+def sint64_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_SINT64, group=group)
 
 
-def float_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_FLOAT)
+def float_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_FLOAT, group=group)
 
 
-def double_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_DOUBLE)
+def double_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_DOUBLE, group=group)
 
 
-def fixed32_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_FIXED32)
+def fixed32_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_FIXED32, group=group)
 
 
-def fixed64_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_FIXED64)
+def fixed64_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_FIXED64, group=group)
 
 
-def sfixed32_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_SFIXED32)
+def sfixed32_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_SFIXED32, group=group)
 
 
-def sfixed64_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_SFIXED64)
+def sfixed64_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_SFIXED64, group=group)
 
 
-def string_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_STRING)
+def string_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_STRING, group=group)
 
 
-def bytes_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_BYTES)
+def bytes_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_BYTES, group=group)
 
 
-def message_field(number: int) -> Any:
-    return dataclass_field(number, TYPE_MESSAGE)
+def message_field(number: int, group: Optional[str] = None) -> Any:
+    return dataclass_field(number, TYPE_MESSAGE, group=group)
 
 
-def map_field(number: int, key_type: str, value_type: str) -> Any:
-    return dataclass_field(number, TYPE_MAP, map_types=(key_type, value_type))
+def map_field(
+    number: int, key_type: str, value_type: str, group: Optional[str] = None
+) -> Any:
+    return dataclass_field(
+        number, TYPE_MAP, map_types=(key_type, value_type), group=group
+    )
 
 
 class Enum(int, enum.Enum):
@@ -383,46 +393,52 @@ class Message(ABC):
     to go between Python, binary and JSON protobuf message representations.
     """
 
-    # True if this message was or should be serialized on the wire. This can
-    # be used to detect presence (e.g. optional wrapper message) and is used
-    # internally during parsing/serialization.
-    serialized_on_wire: bool
-
     def __post_init__(self) -> None:
         # Set a default value for each field in the class after `__init__` has
         # already been run.
+        group_map = {"fields": {}, "groups": {}}
         for field in dataclasses.fields(self):
-            if getattr(self, field.name) != PLACEHOLDER:
-                # Skip anything not set (aka set to the sentinel value)
-                continue
-
             meta = FieldMetadata.get(field)
 
-            t = self._cls_for(field, index=-1)
+            if meta.group:
+                group_map["fields"][field.name] = meta.group
 
-            value: Any = 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)
+                if meta.group not in group_map["groups"]:
+                    group_map["groups"][meta.group] = {"current": None, "fields": set()}
+                group_map["groups"][meta.group]["fields"].add(field)
 
-            setattr(self, field.name, value)
+            if getattr(self, field.name) != PLACEHOLDER:
+                # Skip anything not set to the sentinel value
+
+                if meta.group:
+                    # This was set, so make it the selected value of the one-of.
+                    group_map["groups"][meta.group]["current"] = field
+
+                continue
+
+            setattr(self, field.name, self._get_field_default(field, meta))
 
         # Now that all the defaults are set, reset it!
-        self.__dict__["serialized_on_wire"] = False
+        self.__dict__["_serialized_on_wire"] = False
         self.__dict__["_unknown_fields"] = b""
+        self.__dict__["_group_map"] = group_map
 
     def __setattr__(self, attr: str, value: Any) -> None:
-        if attr != "serialized_on_wire":
+        if attr != "_serialized_on_wire":
             # Track when a field has been set.
-            self.__dict__["serialized_on_wire"] = True
+            self.__dict__["_serialized_on_wire"] = True
+
+        if attr in getattr(self, "_group_map", {}).get("fields", {}):
+            group = self._group_map["fields"][attr]
+            for field in self._group_map["groups"][group]["fields"]:
+                if field.name == attr:
+                    self._group_map["groups"][group]["current"] = field
+                else:
+                    super().__setattr__(
+                        field.name,
+                        self._get_field_default(field, FieldMetadata.get(field)),
+                    )
+
         super().__setattr__(attr, value)
 
     def __bytes__(self) -> bytes:
@@ -434,8 +450,15 @@ class Message(ABC):
             meta = FieldMetadata.get(field)
             value = getattr(self, field.name)
 
+            # Being selected in a a group means this field is the one that is
+            # currently set in a `oneof` group, so it must be serialized even
+            # if the value is the default zero value.
+            selected_in_group = False
+            if meta.group and self._group_map["groups"][meta.group]["current"] == field:
+                selected_in_group = True
+
             if isinstance(value, list):
-                if not len(value):
+                if not len(value) and not selected_in_group:
                     # Empty values are not serialized
                     continue
 
@@ -451,7 +474,7 @@ class Message(ABC):
                     for item in value:
                         output += _serialize_single(meta.number, meta.proto_type, item)
             elif isinstance(value, dict):
-                if not len(value):
+                if not len(value) and not selected_in_group:
                     # Empty values are not serialized
                     continue
 
@@ -461,12 +484,12 @@ 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 == get_default(meta.proto_type):
+                if value == get_default(meta.proto_type) and not selected_in_group:
                     # Default (zero) values are not serialized
                     continue
 
                 serialize_empty = False
-                if isinstance(value, Message) and value.serialized_on_wire:
+                if isinstance(value, Message) and value._serialized_on_wire:
                     serialize_empty = True
                 output += _serialize_single(
                     meta.number, meta.proto_type, value, serialize_empty=serialize_empty
@@ -486,6 +509,24 @@ class Message(ABC):
             cls = type_hints[field.name].__args__[index]
         return cls
 
+    def _get_field_default(self, field: dataclasses.Field, meta: FieldMetadata) -> Any:
+        t = self._cls_for(field, index=-1)
+
+        value: Any = 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)
+
+        return value
+
     def _postprocess_single(
         self, wire_type: int, meta: FieldMetadata, field: dataclasses.Field, value: Any
     ) -> Any:
@@ -508,7 +549,7 @@ class Message(ABC):
             elif meta.proto_type == TYPE_MESSAGE:
                 cls = self._cls_for(field)
                 value = cls().parse(value)
-                value.serialized_on_wire = True
+                value._serialized_on_wire = True
             elif meta.proto_type == TYPE_MAP:
                 # TODO: This is slow, use a cache to make it faster since each
                 #       key/value pair will recreate the class.
@@ -518,8 +559,8 @@ class Message(ABC):
                 Entry = dataclasses.make_dataclass(
                     "Entry",
                     [
-                        ("key", kt, dataclass_field(1, meta.map_types[0], None)),
-                        ("value", vt, dataclass_field(2, meta.map_types[1], None)),
+                        ("key", kt, dataclass_field(1, meta.map_types[0])),
+                        ("value", vt, dataclass_field(2, meta.map_types[1])),
                     ],
                     bases=(Message,),
                 )
@@ -597,7 +638,7 @@ class Message(ABC):
                     # Convert each item.
                     v = [i.to_dict() for i in v]
                     output[field.name] = v
-                elif v.serialized_on_wire:
+                elif v._serialized_on_wire:
                     output[field.name] = v.to_dict()
             elif meta.proto_type == "map":
                 for k in v:
@@ -632,7 +673,7 @@ class Message(ABC):
         Parse the key/value pairs in `value` into this message instance. This
         returns the instance itself and is therefore assignable and chainable.
         """
-        self.serialized_on_wire = True
+        self._serialized_on_wire = True
         for field in dataclasses.fields(self):
             meta = FieldMetadata.get(field)
             if field.name in value and value[field.name] is not None:
@@ -685,6 +726,23 @@ class Message(ABC):
         return self.from_dict(json.loads(value))
 
 
+def serialized_on_wire(message: Message) -> bool:
+    """
+    True if this message was or should be serialized on the wire. This can
+    be used to detect presence (e.g. optional wrapper message) and is used
+    internally during parsing/serialization.
+    """
+    return message._serialized_on_wire
+
+
+def which_one_of(message: Message, group_name: str) -> Tuple[str, Any]:
+    """Return the name and value of a message's one-of field group."""
+    field = message._group_map["groups"].get(group_name, {}).get("current")
+    if not field:
+        return ("", None)
+    return (field.name, getattr(message, field.name))
+
+
 class ServiceStub(ABC):
     """
     Base class for async gRPC service stubs.

betterproto/plugin.py

@@ -242,6 +242,10 @@ def generate_code(request, response):
                             if f.type in [1, 2, 3, 4, 5, 6, 7, 8, 13, 15, 16, 17, 18]:
                                 packed = True
 
+                        one_of = ""
+                        if f.HasField("oneof_index"):
+                            one_of = item.oneof_decl[f.oneof_index].name
+
                         data["properties"].append(
                             {
                                 "name": f.name,
@@ -254,6 +258,7 @@ def generate_code(request, response):
                                 "zero": zero,
                                 "repeated": repeated,
                                 "packed": packed,
+                                "one_of": one_of,
                             }
                         )
                         # print(f, file=sys.stderr)

betterproto/templates/template.py

@@ -44,7 +44,7 @@ class {{ message.name }}(betterproto.Message):
         {% if field.comment %}
 {{ field.comment }}
         {% 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 %})
+    {{ 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 %}{% if field.one_of %}, group="{{ field.one_of }}"{% endif %})
     {% endfor %}
     {% if not message.properties %}
     pass

betterproto/tests/test_features.py

@@ -13,23 +13,23 @@ def test_has_field():
 
     # Unset by default
     foo = Foo()
-    assert foo.bar.serialized_on_wire == False
+    assert betterproto.serialized_on_wire(foo.bar) == False
 
     # Serialized after setting something
     foo.bar.baz = 1
-    assert foo.bar.serialized_on_wire == True
+    assert betterproto.serialized_on_wire(foo.bar) == True
 
     # Still has it after setting the default value
     foo.bar.baz = 0
-    assert foo.bar.serialized_on_wire == True
+    assert betterproto.serialized_on_wire(foo.bar) == True
 
-    # Manual override
-    foo.bar.serialized_on_wire = False
-    assert foo.bar.serialized_on_wire == False
+    # Manual override (don't do this)
+    foo.bar._serialized_on_wire = False
+    assert betterproto.serialized_on_wire(foo.bar) == False
 
     # Can manually set it but defaults to false
     foo.bar = Bar()
-    assert foo.bar.serialized_on_wire == False
+    assert betterproto.serialized_on_wire(foo.bar) == False
 
 
 def test_enum_as_int_json():
@@ -70,3 +70,48 @@ def test_unknown_fields():
 
     new_again = Newer().parse(round_trip)
     assert newer == new_again
+
+
+def test_oneof_support():
+    @dataclass
+    class Sub(betterproto.Message):
+        val: int = betterproto.int32_field(1)
+
+    @dataclass
+    class Foo(betterproto.Message):
+        bar: int = betterproto.int32_field(1, group="group1")
+        baz: str = betterproto.string_field(2, group="group1")
+        sub: Sub = betterproto.message_field(3, group="group2")
+        abc: str = betterproto.string_field(4, group="group2")
+
+    foo = Foo()
+
+    assert betterproto.which_one_of(foo, "group1")[0] == ""
+
+    foo.bar = 1
+    foo.baz = "test"
+
+    # Other oneof fields should now be unset
+    assert foo.bar == 0
+    assert betterproto.which_one_of(foo, "group1")[0] == "baz"
+
+    foo.sub.val = 1
+    assert betterproto.serialized_on_wire(foo.sub)
+
+    foo.abc = "test"
+
+    # Group 1 shouldn't be touched, group 2 should have reset
+    assert foo.sub.val == 0
+    assert betterproto.serialized_on_wire(foo.sub) == False
+    assert betterproto.which_one_of(foo, "group2")[0] == "abc"
+
+    # Zero value should always serialize for one-of
+    foo = Foo(bar=0)
+    assert betterproto.which_one_of(foo, "group1")[0] == "bar"
+    assert bytes(foo) == b"\x08\x00"
+
+    # Round trip should also work
+    foo2 = Foo().parse(bytes(foo))
+    assert betterproto.which_one_of(foo2, "group1")[0] == "bar"
+    assert foo.bar == 0
+    assert betterproto.which_one_of(foo2, "group2")[0] == ""