Repository: caikit/caikit
Pull request #132 (2023-05-23 16:28:15)

🐛 handle from_proto oneof case
## Issue

The problem is the fact that when we have
```
    @dataobject
    class BazObj(DataObjectBase):
        @dataobject
        class Foo(DataObjectBase):
            data: List[str]

        @dataobject
        class Bar(DataObjectBase):
            data: str

        data_stream: Union[
            Annotated[Foo, FieldNumber(1), OneofField("foo")],
            Annotated[Bar, FieldNumber(2), OneofField("bar")],
        ]
```

During the `from_proto` -> since `foo` and `bar` are part of `cls._fields_message`, we do the `if proto.HasField(field): `check before setting the value of the field. So in this case, we don’t set the value of `bar`.

But, in the case of
```
    @dataobject
    class Foo(DataObjectBase):
        foo: Union[
            Annotated[int, FieldNumber(10), OneofField("foo_int")],
            Annotated[float, FieldNumber(20), OneofField("foo_float")],
        ]
```
since both `foo_int` and `foo_float` are in `cls._fields_primitive`, we don’t check for `hasField` and instead assign the field value directly, and so we set both `foo_int` and `foo_float`.
This crashes the `_make_oneof_init` within `dataobject` which wants only 1 field within a `oneof` to be set at a time.


## Possible solution
I figured the only way to correct this behavior is to tackle it at the `proto` level - because once we are beyond the `proto` level, there’s no way for us to figure out an unset field vs false value of the field.


Patch:
diff --git a/caikit/core/data_model/base.py b/caikit/core/data_model/base.py
index 4c3fc6ff..c1096dbd 100644
--- a/caikit/core/data_model/base.py
+++ b/caikit/core/data_model/base.py
@@ -597,8 +597,9 @@ class DataBase(metaclass=_DataBaseMetaClass):
                 )
 
             if field in cls._fields_primitive or field in cls._fields_enum:
-                kwargs[field] = proto_attr
-
+                # special case for oneofs
+                if field not in cls._fields_to_oneof or proto.HasField(field):
+                    kwargs[field] = proto_attr
             elif (
                 field in cls._fields_primitive_repeated
                 or field in cls._fields_enum_repeated


Test patch:
diff --git a/tests/core/data_model/test_dataobject.py b/tests/core/data_model/test_dataobject.py
index 5d04c1ea..dc7cc17d 100644
--- a/tests/core/data_model/test_dataobject.py
+++ b/tests/core/data_model/test_dataobject.py
@@ -459,6 +459,34 @@ def test_dataobject_with_same_type_of_oneof():
     assert foo2.foo_bool2
 
 
+def test_dataobject_primitive_oneof_round_trips():
+    @dataobject
+    class Foo(DataObjectBase):
+        foo: Union[
+            Annotated[int, FieldNumber(10), OneofField("foo_int")],
+            Annotated[float, FieldNumber(20), OneofField("foo_float")],
+        ]
+
+    # proto round trip
+    foo1 = Foo(foo_int=2)
+    assert foo1.which_oneof("foo") == "foo_int"
+    proto_repr_foo = foo1.to_proto()
+    assert Foo.from_proto(proto=proto_repr_foo).to_proto() == proto_repr_foo
+
+    foo2 = Foo(foo=2)
+    assert foo2.which_oneof("foo") == "foo_int"
+    proto_repr_foo = foo2.to_proto()
+    assert Foo.from_proto(proto=proto_repr_foo).to_proto() == proto_repr_foo
+
+    # json round trip
+    json_repr_foo = foo1.to_json()
+    assert json.loads(json_repr_foo) == {
+        "foo_int": 2,
+        "foo_float": None,
+    }
+    assert Foo.from_json(json_repr_foo) == foo1
+
+
 def test_dataobject_round_trip_json():
     """Make sure that a dataobject class can serialize to/from json"""