Repository: caikit/caikit
Pull request #143 (2023-05-24 23:56:14)

♻️ ignore unset fields in to_dict
```
class Foo(DataObjectBase):
    foo: Union[
        Annotated[int, FieldNumber(10), OneofField("foo_int")],
        Annotated[float, FieldNumber(20), OneofField("foo_float")],
    ]
```

```
foo1 = Foo(foo_int=2)
json_repr_foo = foo1.to_json()
assert json.loads(json_repr_foo) == {
    "foo_int": 2,
    "foo_float": None,
}
```

`"foo_float": None` should not be in the `to_json` representation.


Patch:
diff --git a/caikit/core/data_model/base.py b/caikit/core/data_model/base.py
index c1096dbd..471dc519 100644
--- a/caikit/core/data_model/base.py
+++ b/caikit/core/data_model/base.py
@@ -267,7 +267,6 @@ class _DataBaseMetaClass(type):
         # existed in old_fields for supporting oneofs
         # see https://github.com/caikit/caikit/pull/107 for details
         for field in set(cls.fields + tuple(old_fields)):
-
             # If the field is the name of a field within a oneof and it was not
             # in the old fields, the data is held under the oneof's name if this
             # is the set value for the oneof
@@ -371,7 +370,6 @@ class _DataBaseMetaClass(type):
 
             if num_kwargs > 0:  # Do a quick check for performance reason
                 for field_name, field_val in kwargs.items():
-
                     # If this is a oneof field, alias to the oneof name
                     if oneof_name := cls._fields_to_oneof.get(field_name):
                         which_oneof[oneof_name] = field_name
@@ -775,7 +773,15 @@ class DataBase(metaclass=_DataBaseMetaClass):
 
     def to_dict(self) -> dict:
         """Convert to a dictionary representation."""
-        return {field: self._field_to_dict_element(field) for field in self.fields}
+        # maintain a list of fields to convert to dict, special handling for oneofs
+        fields_to_dict = []
+        for field in self.fields:
+            if (
+                not field in self._fields_to_oneof
+                or self.which_oneof(self._fields_to_oneof[field]) == field
+            ):
+                fields_to_dict.append(field)
+        return {field: self._field_to_dict_element(field) for field in fields_to_dict}
 
     def to_json(self, **kwargs) -> str:
         """Convert to a json representation."""


Test patch:
diff --git a/tests/core/data_model/test_dataobject.py b/tests/core/data_model/test_dataobject.py
index dc7cc17d..9c10198a 100644
--- a/tests/core/data_model/test_dataobject.py
+++ b/tests/core/data_model/test_dataobject.py
@@ -381,6 +381,13 @@ def test_dataobject_with_oneof():
     assert bar1.bar is bar1.data_stream
     assert bar1.foo is None
 
+    # Test to_dict
+    dict_repr_foo = foo1.to_dict()
+    assert dict_repr_foo == {
+        "foo": {"data": ["hello"]},
+    }
+    assert dict_repr_foo["foo"]["data"] == ["hello"]
+
     # Test proto round trip
     proto_repr_foo = foo1.to_proto()
     assert proto_repr_foo.foo.data == ["hello"]
@@ -392,7 +399,6 @@ def test_dataobject_with_oneof():
     json_repr_foo = foo1.to_json()
     assert json.loads(json_repr_foo) == {
         "foo": {"data": ["hello"]},
-        "bar": None,
     }
     assert BazObj.from_json(json_repr_foo) == foo1
 
@@ -478,11 +484,13 @@ def test_dataobject_primitive_oneof_round_trips():
     proto_repr_foo = foo2.to_proto()
     assert Foo.from_proto(proto=proto_repr_foo).to_proto() == proto_repr_foo
 
+    # dict test
+    assert foo1.to_dict() == {"foo_int": 2}
+
     # 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