Repository: caikit/caikit
Pull request #122 (2023-05-17 22:29:09)

Full oneof support
## Description

This PR adds full support for `oneof` semantics to the data model. Specifically, it makes the following true:

1. Any class derived from `DataBase` (including those derived from `DataObjectBase`) will internally store the current value of a oneof under the name of the oneof and not the name of the currently set field within the oneof
2. The `__init__` to any class derived from `DataBase` can take a keyword argument with the name of the `oneof` OR the name of one of the fields, but not both (and not multiple fields)
3. When calling `getattr` on a data model instance with the name of a `oneof`, it will return whatever value is set if any is
4. When calling `getattr` on a data model instance with the name of a `oneof` field, it will return the value of the oneof if (and only if) that field is the one set in the oneof
5. Every object has a `which_oneof` function which can be used to determine which field in a oneof is currently set (if any)
6. When the value is set using the name of the `oneof` rather than the name of a field, the `which_oneof` information is lazily introspected from the type of the value the first time `which_oneof` is called (this includes attribute access by field name)
    * **NOTE**: This may cause issues with `oneofs` where the fields are not disambiguated by type! In this case, the first matching field is chosen 


Patch:
diff --git a/caikit/core/data_model/base.py b/caikit/core/data_model/base.py
index 8dc48d75..4c3fc6ff 100644
--- a/caikit/core/data_model/base.py
+++ b/caikit/core/data_model/base.py
@@ -21,13 +21,14 @@
 
 # Standard
 from enum import Enum
-from typing import Optional, Type, Union
+from typing import Any, Dict, Optional, Type, Union
 import base64
 import json
 
 # Third Party
 from google.protobuf import json_format
 from google.protobuf.descriptor import Descriptor
+from google.protobuf.internal import type_checkers as proto_type_checkers
 from google.protobuf.message import Message as ProtoMessageType
 
 # First Party
@@ -59,6 +60,11 @@ class _DataBaseMetaClass(type):
     # construction.
     _FWD_DECL_FIELDS = "__fwd_decl_fields__"
 
+    # Special instance attributes that an instance of a class derived from
+    # DataBase may have. These are added to __slots__.
+    _BACKEND_ATTR = "_backend"
+    _WHICH_ONEOF_ATTR = "_which_oneof"
+
     def __new__(mcs, name, bases, attrs):
         """When constructing a new data model class, we set the 'fields' class variable from the
         protobufs descriptor and then set the '__slots__' magic class attribute to fields.  This
@@ -70,29 +76,43 @@ class _DataBaseMetaClass(type):
         protobufs, it can be named in the tuple class variable _private_slots and will
         automatically be added to __slots__.
         """
-        # Get all fields in protobufs with same name as class,
-        # except for DataBase, which has no matching protobufs
-        fields = ()
 
-        # Protobufs fields can be divided into these categories, which are used to automatically
-        # determine appropriate behavior in a number of methods
-        fields_enum_map = {}
-        fields_enum_rev = {}
-        _fields_message = ()
-        _fields_message_repeated = ()
-        _fields_enum = ()
-        _fields_enum_repeated = ()
-        _fields_primitive = ()
-        _fields_primitive_repeated = ()
-        _fields_map = ()
+        # Protobufs fields can be divided into these categories, which are used
+        # to automatically determine appropriate behavior in a number of methods
+        attrs["full_name"] = name
+        attrs["fields_enum_map"] = {}
+        attrs["fields_enum_rev"] = {}
+        attrs["_fields_oneofs_map"] = {}
+        attrs["_fields_to_oneof"] = {}
+        attrs["_fields_map"] = ()
+        attrs["_fields_message"] = ()
+        attrs["_fields_message_repeated"] = ()
+        attrs["_fields_enum"] = ()
+        attrs["_fields_enum_repeated"] = ()
+        attrs["_fields_primitive"] = ()
+        attrs["_fields_primitive_repeated"] = ()
+
+        # Look for the set of fields either from a predefined protobuf class or
+        # from a forward declaration from @dataobject
+        fields = ()
         proto_class = None
-        full_name = name
         if name not in ["DataBase", "DataObjectBase"]:
             # Look for a precompiled proto class and if found, parse its
             # descriptor
             proto_class = attrs.get("_proto_class")
             if proto_class is not None:
-                fields = tuple(proto_class.DESCRIPTOR.fields_by_name)
+                all_oneof_fields = [
+                    field.name
+                    for oneof in proto_class.DESCRIPTOR.oneofs
+                    for field in oneof.fields
+                ]
+                fields = tuple(
+                    (
+                        field
+                        for field in proto_class.DESCRIPTOR.fields_by_name
+                        if field not in all_oneof_fields
+                    )
+                ) + tuple(proto_class.DESCRIPTOR.oneofs_by_name)
 
             # Otherwise, we need to get the fields from a "special" attribute
             else:
@@ -106,6 +126,8 @@ class _DataBaseMetaClass(type):
                     "No proto class found for {}",
                     name,
                 )
+        attrs["fields"] = fields
+        attrs["_proto_class"] = proto_class
 
         # Look if any private slots are declared as class variables
         private_slots = attrs.setdefault("_private_slots", ())
@@ -113,25 +135,12 @@ class _DataBaseMetaClass(type):
         # Class slots are fields + private slots, this prevents other
         # member attributes from being set and also improves performance
         attrs["__slots__"] = tuple(
-            [f"_{field}" for field in fields] + list(private_slots) + ["_backend"]
+            [f"_{field}" for field in fields]
+            + list(private_slots)
+            + [mcs._BACKEND_ATTR, mcs._WHICH_ONEOF_ATTR]
         )
 
-        # Set fields class variable for reference
-        # these are valuable for validating attributes and
-        # also for recursively converting to and from protobufs
-        attrs["full_name"] = full_name
-        attrs["fields"] = fields
-        attrs["fields_enum_map"] = fields_enum_map
-        attrs["fields_enum_rev"] = fields_enum_rev
-        attrs["_fields_map"] = _fields_map
-        attrs["_fields_message"] = _fields_message
-        attrs["_fields_message_repeated"] = _fields_message_repeated
-        attrs["_fields_enum"] = _fields_enum
-        attrs["_fields_enum_repeated"] = _fields_enum_repeated
-        attrs["_fields_primitive"] = _fields_primitive
-        attrs["_fields_primitive_repeated"] = _fields_primitive_repeated
-        attrs["_proto_class"] = proto_class
-
+        # Create the instance of the type
         instance = super().__new__(mcs, name, bases, attrs)
 
         # If there's a valid proto class, perform proto descriptor parsing
@@ -154,6 +163,7 @@ class _DataBaseMetaClass(type):
 
         # preserve old fields for _make_property_getter later
         old_fields = cls.fields
+
         # overwrite to only have proto-specific fields present
         cls.fields = tuple(cls._proto_class.DESCRIPTOR.fields_by_name)
 
@@ -171,6 +181,21 @@ class _DataBaseMetaClass(type):
             if field.enum_type is not None
         }
 
+        # mapping of all oneofs and the fields that are part of them
+        # NOTE: protobuf makes an interesting use of oneof to wrap types that
+        #   should be explicitly optional. We don't want to consider these
+        #   oneofs in the general oneof handling.
+        cls._fields_oneofs_map = {
+            oneof_name: [field.name for field in oneof.fields]
+            for oneof_name, oneof in cls._proto_class.DESCRIPTOR.oneofs_by_name.items()
+            if len(oneof.fields) != 1 or oneof.name != f"_{oneof.fields[0].name}"
+        }
+        cls._fields_to_oneof = {
+            field_name: oneof_name
+            for (oneof_name, oneof_fields) in cls._fields_oneofs_map.items()
+            for field_name in oneof_fields
+        }
+
         # all repeated fields
         fields_repeated = tuple(
             field.name
@@ -215,8 +240,10 @@ class _DataBaseMetaClass(type):
             _fields_message_all
         )
 
+        # enums that are not repeated
         cls._fields_enum = frozenset(_fields_enum_all).difference(fields_repeated)
 
+        # enums that are repeated
         cls._fields_enum_repeated = frozenset(_fields_enum_all).intersection(
             fields_repeated
         )
@@ -236,12 +263,21 @@ class _DataBaseMetaClass(type):
         #   registry shared for all
         _DataBaseMetaClass.class_registry[cls.full_name] = cls
 
-        # Add properties that use the underlying backend
-        # also add fields that existed in old_fields
-        # for supporting oneofs
+        # Add properties that use the underlying backend. Also add fields that
+        # 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)):
-            setattr(cls, field, mcs._make_property_getter(field))
+
+            # 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
+            if oneof_name := cls._fields_to_oneof.get(field):
+                setattr(cls, field, mcs._make_property_getter(field, oneof_name))
+
+            # If the field is a plain field or the name of a oneof, it will be
+            # accessed directly
+            else:
+                setattr(cls, field, mcs._make_property_getter(field))
 
         # If there is not already an __init__ function defined, make one
         current_init = cls.__init__
@@ -249,15 +285,14 @@ class _DataBaseMetaClass(type):
             setattr(cls, "__init__", mcs._make_init(cls.fields))
 
     @classmethod
-    def _make_property_getter(mcs, field):
+    def _make_property_getter(mcs, field, oneof_name=None):
         """This helper creates an @property attribute getter for the given field
 
         NOTE: This needs to live as a standalone function in order for the given
             field name to be properly bound to the closure for the attrs
         """
-        private_name = f"_{field}"
+        private_name = f"_{field}" if oneof_name is None else oneof_name
 
-        @property
         def _property_getter(self):
             # Check to see if the private name is defined and just return it if
             # it is
@@ -266,17 +301,35 @@ class _DataBaseMetaClass(type):
                 return current
 
             # If not currently set, delegate to the backend
-            attr_val = self._backend.get_attribute(self.__class__, field)
+            backend = self.backend
+            if backend is None:
+                error(
+                    "<COR66616239E>",
+                    AttributeError(
+                        f"{type(self)} missing attribute {field} and no backend set"
+                    ),
+                )
+            attr_val = backend.get_attribute(self.__class__, field)
 
             # If the backend says that this attribute should be cached, set it
             # as an attribute on the class
-            if self._backend.cache_attribute(field, attr_val):
-                setattr(self, private_name, attr_val)
+            if backend.cache_attribute(field, attr_val):
+                setattr(self, field, attr_val)
 
             # Return the value found by the backend
             return attr_val
 
-        return _property_getter
+        # If this is a oneof, add an extra layer of wrapping to check
+        # which_oneof before returning a valid result
+        if oneof_name:
+
+            def _oneof_property_getter(self):
+                if self.which_oneof(oneof_name) == field:
+                    return _property_getter(self)
+
+            return property(_oneof_property_getter)
+
+        return property(_property_getter)
 
     @staticmethod
     def _make_init(fields):
@@ -298,6 +351,12 @@ class _DataBaseMetaClass(type):
             num_fields = len(fields)
             used_fields = []
 
+            # If the proto has oneofs, set up which_oneof
+            which_oneof = {}
+            cls = self.__class__
+            if cls._fields_oneofs_map:
+                setattr(self, _DataBaseMetaClass._WHICH_ONEOF_ATTR, which_oneof)
+
             if num_args + num_kwargs > num_fields:
                 error(
                     "<COR71444420E>",
@@ -312,7 +371,16 @@ class _DataBaseMetaClass(type):
 
             if num_kwargs > 0:  # Do a quick check for performance reason
                 for field_name, field_val in kwargs.items():
-                    if field_name not in fields:
+
+                    # 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
+                        field_name = oneof_name
+
+                    if (
+                        field_name not in fields
+                        and field_name not in cls._fields_oneofs_map
+                    ):
                         error(
                             "<COR71444421E>", TypeError(f"Unknown field {field_name}")
                         )
@@ -327,7 +395,10 @@ class _DataBaseMetaClass(type):
             # Default all unspecified fields to None
             if num_fields > 0:  # Do a quick check for performance reason
                 for field_name in fields:
-                    if field_name not in used_fields:
+                    if (
+                        field_name not in used_fields
+                        and field_name not in cls._fields_to_oneof
+                    ):
                         setattr(self, field_name, None)
 
         # Set docstring to the method explicitly
@@ -344,10 +415,22 @@ class DataBase(metaclass=_DataBaseMetaClass):
     """
 
     def __setattr__(self, name, val):
-        """If attempting to set one of the named fields, instead set the
-        "private" version of the attribute.
+        """Handle attribute setting for oneofs and named fields with delegation
+        to backends as needed
         """
-        if name in self.__class__.fields:
+        # If setting a oneof directly, remove any oneof information
+        cls = self.__class__
+        if name in cls._fields_oneofs_map:
+            self._get_which_oneof_dict().pop(name, None)
+
+        # If this is the name of a oneof field, set the oneof itself
+        if oneof_name := cls._fields_to_oneof.get(name):
+            self._get_which_oneof_dict()[oneof_name] = name
+            name = oneof_name
+
+        # If attempting to set one of the named fields or a oneof, instead set
+        # the private version of the attribute.
+        if name in cls.fields or name in cls._fields_oneofs_map:
             super().__setattr__(f"_{name}", val)
         else:
             super().__setattr__(name, val)
@@ -383,9 +466,85 @@ class DataBase(metaclass=_DataBaseMetaClass):
     @classmethod
     def from_backend(cls, backend):
         instance = cls.__new__(cls)
-        setattr(instance, "_backend", backend)
+        setattr(instance, _DataBaseMetaClass._BACKEND_ATTR, backend)
         return instance
 
+    @property
+    def backend(self) -> Optional["DataModelBackendBase"]:
+        return getattr(self, _DataBaseMetaClass._BACKEND_ATTR, None)
+
+    def which_oneof(self, oneof_name: str) -> Optional[str]:
+        """Get the name of the oneof field set for the given oneof or None if no
+        field is set
+        """
+        # If the internal dict is already set, use that information
+        which_oneof = self._get_which_oneof_dict()
+        if current_val := which_oneof.get(oneof_name):
+            return current_val
+
+        # Get the current value for the oneof and introspect which field its
+        # type matches
+        oneof_val = getattr(self, oneof_name)
+        which_field = self._infer_which_oneof(oneof_name, oneof_val)
+        if which_field is not None:
+            which_oneof[oneof_name] = which_field
+        return which_field
+
+    @classmethod
+    def _infer_which_oneof(cls, oneof_name: str, oneof_val: Any) -> Optional[str]:
+        """Check each candidate field within the oneof to see if it's a type
+        match
+
+        NOTE: In the case where fields within a oneof have the same type, the
+          first field whose type matches will be used!
+        """
+        for field_name in cls._fields_oneofs_map.get(oneof_name, []):
+            if cls._is_valid_type_for_field(field_name, oneof_val):
+                return field_name
+
+    def _get_which_oneof_dict(self) -> Dict[str, str]:
+        which_oneof = getattr(self, _DataBaseMetaClass._WHICH_ONEOF_ATTR, None)
+        if which_oneof is None:
+            super().__setattr__(_DataBaseMetaClass._WHICH_ONEOF_ATTR, {})
+            which_oneof = getattr(self, _DataBaseMetaClass._WHICH_ONEOF_ATTR)
+        return which_oneof
+
+    @classmethod
+    def _is_valid_type_for_field(cls, field_name: str, val: Any) -> bool:
+        """Check whether the given value is valid for the given field"""
+        field_descriptor = cls._proto_class.DESCRIPTOR.fields_by_name[field_name]
+
+        if val is None:
+            return False
+
+        # If it's a data object or an enum and the descriptors match, it's a
+        # good type
+        if (
+            isinstance(val, DataBase)
+            and field_descriptor.message_type == val.get_proto_class().DESCRIPTOR
+        ) or (
+            isinstance(val, Enum)
+            and field_descriptor.enum_type == val.get_proto_class().DESCRIPTOR
+        ):
+            return True
+
+        # If it's a data object or an enum and the descriptors don't match, it's
+        # a bad type
+        if field_descriptor.type in [
+            field_descriptor.TYPE_MESSAGE,
+            field_descriptor.TYPE_ENUM,
+        ]:
+            return False
+
+        # If it's a primitive, use protobuf type checkers
+        checker = proto_type_checkers.GetTypeChecker(field_descriptor)
+        try:
+            checker.CheckValue(val)
+            return True
+        except TypeError:
+            pass
+        return False
+
     @classmethod
     def from_binary_buffer(cls, buf):
         """Builds the data model object out of the binary string
@@ -413,12 +572,12 @@ class DataBase(metaclass=_DataBaseMetaClass):
             protobufs
                 A DataBase object.
         """
-        if cls.__name__ != proto.DESCRIPTOR.name:
+        if cls._proto_class.DESCRIPTOR.name != proto.DESCRIPTOR.name:
             error(
                 "<COR71783894E>",
                 ValueError(
                     "class name `{}` does not match protobufs name `{}`".format(
-                        cls.__name__, proto.DESCRIPTOR.name
+                        cls._proto_class.DESCRIPTOR.name, proto.DESCRIPTOR.name
                     )
                 ),
             )
diff --git a/caikit/core/data_model/data_backends/dict_backend.py b/caikit/core/data_model/data_backends/dict_backend.py
index 250cc6d2..bd531b54 100644
--- a/caikit/core/data_model/data_backends/dict_backend.py
+++ b/caikit/core/data_model/data_backends/dict_backend.py
@@ -63,7 +63,10 @@ class DictBackend(DataModelBackendBase):
         #   being used correctly.
 
         # Make sure the name is a valid field on the given class
-        if name not in data_model_class.fields:
+        if (
+            name not in data_model_class.fields
+            and name not in data_model_class._fields_oneofs_map
+        ):
             error(
                 "<COR85037211E>",
                 AttributeError(
diff --git a/caikit/core/data_model/dataobject.py b/caikit/core/data_model/dataobject.py
index 32685e33..75e00017 100644
--- a/caikit/core/data_model/dataobject.py
+++ b/caikit/core/data_model/dataobject.py
@@ -67,16 +67,11 @@ class _DataObjectBaseMetaClass(_DataBaseMetaClass):
 
         # Get the annotations that will go into the dataclass
         if name != "DataObjectBase":
-            raw_field_names = attrs.get("__annotations__")
-            if raw_field_names is None:
+            field_names = attrs.get("__annotations__")
+            if field_names is None:
                 raise TypeError(
                     "All DataObjectBase classes must follow dataclass syntax"
                 )
-
-            # TODO: Sort out oneof field names
-            field_names = list(raw_field_names.keys())
-
-            # Add the forward declaration
             attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names
 
         # Delegate to the base metaclass
@@ -129,47 +124,27 @@ def dataobject(*args, **kwargs) -> Callable[[Type], Type[DataBase]]:
         # Add the package to the kwargs
         kwargs.setdefault("package", package)
 
-        # If there's a schema in the keyword args, use jtd_to_proto
-        schema = kwargs.pop("schema", None)
-        if schema is not None:
-            log.debug("Using JTD To Proto")
-            kwargs.setdefault("validate_jtd", True)
-            descriptor = py_to_proto.jtd_to_proto(
-                name=cls.__name__,
-                jtd_def=schema,
-                **kwargs,
-            )
-        # If it's already a dataclass, convert it directly
-        else:
-            log.debug("Using dataclass/enum to proto on dataclass")
-
-            # If it's not an enum, fill in any missing field defaults as None
-            # and make sure it's a dataclass
-            if not issubclass(cls, Enum):
-                log.debug2("Wrapping data class %s", cls)
-                user_defined_defaults = {}
-                for annotation in getattr(cls, "__annotations__", {}):
-                    user_defined_default = getattr(cls, annotation, dataclasses.MISSING)
-                    if user_defined_default == dataclasses.MISSING:
-                        log.debug3("Filling in None default for %s.%s", cls, annotation)
-                        setattr(cls, annotation, None)
-                    else:
-                        user_defined_defaults[annotation] = user_defined_default
-                # If the current __init__ is auto-generated by dataclass, remove
-                # it so that a new one is created with the new defaults. This is
-                # a little hard to detect across different python versions, so
-                # the most reliable way is to assume that the only place
-                # __annotations__ are added to the __init__ function itself are
-                # in dataclass. If cls is either not a dataclass or is a
-                # dataclass with a non-default __init__, there will not be
-                # annotations
-                if getattr(cls.__init__, "__annotations__", None):
-                    log.debug3("Resetting default dataclass init")
-                    delattr(cls, "__init__")
-                cls = dataclasses.dataclass(cls)
-                setattr(cls, _USER_DEFINED_DEFAULTS, user_defined_defaults)
-
-            descriptor = _dataobject_to_proto(dataclass_=cls, **kwargs)
+        # If it's not an enum, fill in any missing field defaults as None
+        # and make sure it's a dataclass
+        if not issubclass(cls, Enum):
+            log.debug2("Wrapping data class %s", cls)
+            user_defined_defaults = {}
+            for annotation in getattr(cls, "__annotations__", {}):
+                user_defined_default = getattr(cls, annotation, dataclasses.MISSING)
+                if user_defined_default == dataclasses.MISSING:
+                    log.debug3("Filling in None default for %s.%s", cls, annotation)
+                    setattr(cls, annotation, None)
+                else:
+                    user_defined_defaults[annotation] = user_defined_default
+            # If the current __init__ is auto-generated by dataclass, remove
+            # it so that a new one is created with the new defaults
+            if _has_dataclass_init(cls):
+                log.debug3("Resetting default dataclass init")
+                delattr(cls, "__init__")
+            cls = dataclasses.dataclass(cls)
+            setattr(cls, _USER_DEFINED_DEFAULTS, user_defined_defaults)
+
+        descriptor = _dataobject_to_proto(dataclass_=cls, **kwargs)
 
         # Create the message class from the dataclass
         proto_class = py_to_proto.descriptor_to_message_class(descriptor)
@@ -185,6 +160,13 @@ def dataobject(*args, **kwargs) -> Callable[[Type], Type[DataBase]]:
         if isinstance(proto_class, type):
             setattr(cls, "_proto_class", proto_class)
             cls = _make_data_model_class(proto_class, cls)
+
+            # If this was a default-generated dataclass __init__ and there are
+            # any oneofs, we need to augment the __init__ to support kwargs for
+            # the individual fields
+            if _has_dataclass_init(cls) and cls._fields_oneofs_map:
+                setattr(cls, "__init__", _make_oneof_init(cls))
+
         else:
             enums.import_enum(proto_class, cls)
             setattr(cls, "_proto_enum", proto_class)
@@ -321,3 +303,65 @@ def _make_data_model_class(proto_class, cls):
         )
 
     return cls
+
+
+def _make_oneof_init(cls):
+    """Helper to augment a defaulted dataclass __init__ to support kwargs for
+    oneof fields
+    """
+    original_init = cls.__init__
+    fields_to_oneofs = cls._fields_to_oneof
+    oneofs_to_fields = cls._fields_oneofs_map
+
+    def __init__(self, *args, **kwargs):
+        new_kwargs = {}
+        to_remove = []
+        which_oneof = {}
+        for field_name, val in kwargs.items():
+            if oneof_name := fields_to_oneofs.get(field_name):
+                oneof_pos_idx = list(cls.__dataclass_fields__.keys()).index(oneof_name)
+                has_pos_val = len(args) > oneof_pos_idx
+                if has_pos_val:
+                    error(
+                        "<COR09282193E>",
+                        TypeError(
+                            "Received conflicting oneof args/kwargs for {}/{}".format(
+                                oneof_name,
+                                field_name,
+                            )
+                        ),
+                    )
+
+                other_oneof_fields = (
+                    field
+                    for field in [oneof_name] + oneofs_to_fields[oneof_name]
+                    if field != field_name
+                )
+                if any(field in kwargs for field in other_oneof_fields):
+                    error(
+                        "<COR59933157E>",
+                        TypeError(
+                            "Received multiple keyword arguments for oneof {}".format(
+                                oneof_name,
+                            )
+                        ),
+                    )
+                new_kwargs[oneof_name] = val
+                to_remove.append(field_name)
+                which_oneof[oneof_name] = field_name
+
+        for kwarg in to_remove:
+            del kwargs[kwarg]
+        kwargs.update(new_kwargs)
+        original_init(self, *args, **kwargs)
+        setattr(self, _DataBaseMetaClass._WHICH_ONEOF_ATTR, which_oneof)
+
+    return __init__
+
+
+def _has_dataclass_init(cls) -> bool:
+    """When the dataclass decorator adds an __init__ to a class, it adds
+    __annotations__ to the init function itself. This function uses that fact to
+    detect if the class's __init__ function was generated by @dataclass
+    """
+    return bool(getattr(cls.__init__, "__annotations__", None))
diff --git a/caikit/runtime/service_generation/data_stream_source.py b/caikit/runtime/service_generation/data_stream_source.py
index 579e3405..38c7fe19 100644
--- a/caikit/runtime/service_generation/data_stream_source.py
+++ b/caikit/runtime/service_generation/data_stream_source.py
@@ -15,7 +15,7 @@
 # Standard
 from datetime import datetime
 from glob import glob
-from typing import Any, Optional, Type
+from typing import Any, List, Optional, Type, Union
 import os
 import sys
 
@@ -24,6 +24,7 @@ from google.protobuf.message import Message as ProtoMessageType
 import grpc
 
 # First Party
+from py_to_proto.dataclass_to_proto import Annotated, OneofField
 import alog
 
 # Local
@@ -31,6 +32,7 @@ from caikit.core.data_model.base import DataBase
 from caikit.core.data_model.dataobject import (
     DataObjectBase,
     _DataObjectBaseMetaClass,
+    _make_oneof_init,
     dataobject,
 )
 from caikit.core.data_model.streams.data_stream import DataStream
@@ -63,16 +65,6 @@ class DataStreamSourceBase(DataStream):
     def __init__(self):
         """Validate oneof semantics"""
         super().__init__(lambda: self.to_data_stream().generator_func())
-        self._set_fields = {
-            f.name: getattr(self, f.name)
-            for f in self.get_proto_class().DESCRIPTOR.oneofs[0].fields
-            if getattr(self, f.name) is not None
-        }
-        if len(self._set_fields) > 1:
-            raise CaikitRuntimeException(
-                grpc.StatusCode.INVALID_ARGUMENT,
-                f"Multiple oneof fields set: {self._set_fields.keys()}",
-            )
 
     def __getstate__(self) -> bytes:
         """A DataStreamSource is pickled by serializing its source
@@ -88,11 +80,7 @@ class DataStreamSourceBase(DataStream):
         oneof members.
         """
         new_inst = self.__class__.from_binary_buffer(pickle_bytes)
-        new_inst_dict = new_inst.__dict__
-        self.__dict__.update(new_inst_dict)
-        for field_descriptor in self.get_proto_class().DESCRIPTOR.oneofs[0].fields:
-            field_name = field_descriptor.name
-            setattr(self, field_name, getattr(new_inst, field_name))
+        setattr(self, new_inst.which_oneof("data_stream"), new_inst.data_stream)
 
     # pylint: disable=too-many-return-statements
     def to_data_stream(self) -> DataStream:
@@ -216,65 +204,92 @@ def make_data_stream_source(data_element_type: Type) -> Type[DataBase]:
     if data_element_type not in _DATA_STREAM_SOURCE_TYPES:
         cls_name = _make_data_stream_source_type_name(data_element_type)
         log.debug("Creating DataStreamSource[%s] -> %s", data_element_type, cls_name)
-        element_type = (
-            data_element_type.get_proto_class().DESCRIPTOR
-            if isinstance(data_element_type, type)
-            and issubclass(data_element_type, DataBase)
-            else _NATIVE_TYPE_TO_JTD[data_element_type]
+
+        # Set up the "sub classes." In python, this is the same as creating a
+        # standalone class with a __qualname__ that nests it under a parent
+        # class. We do this outside the declaration of the parent class so that
+        # these classes can be referenced within the Union annotation for the
+        # outer class itself.
+        package = "caikit_data_model.runtime"
+        JsonData = dataobject(package=package, name=f"{cls_name}JsonData")(
+            _DataObjectBaseMetaClass.__new__(
+                _DataObjectBaseMetaClass,
+                name="JsonData",
+                bases=(DataObjectBase,),
+                attrs={
+                    "__annotations__": {"data": List[data_element_type]},
+                    "__qualname__": f"{cls_name}.JsonData",
+                },
+            ),
         )
-        schema = {
-            "properties": {
-                "data_stream": {
-                    "discriminator": "data_reference_type",
-                    "mapping": {
-                        "JsonData": {
-                            "properties": {
-                                "data": {
-                                    "elements": {"type": element_type},
-                                },
-                            },
-                        },
-                        "File": {"properties": {"filename": {"type": "string"}}},
-                        "ListOfFiles": {
-                            "properties": {
-                                "files": {
-                                    "elements": {"type": "string"},
-                                },
-                            },
-                        },
-                        "Directory": {
-                            "properties": {
-                                "dirname": {"type": "string"},
-                                "extension": {"type": "string"},
-                            }
-                        },
+        File = dataobject(package=package, name=f"{cls_name}File")(
+            _DataObjectBaseMetaClass.__new__(
+                _DataObjectBaseMetaClass,
+                name="File",
+                bases=(DataObjectBase,),
+                attrs={
+                    "__annotations__": {"filename": str},
+                    "__qualname__": f"{cls_name}.File",
+                },
+            ),
+        )
+        ListOfFiles = dataobject(package=package, name=f"{cls_name}ListOfFiles")(
+            _DataObjectBaseMetaClass.__new__(
+                _DataObjectBaseMetaClass,
+                name="ListOfFiles",
+                bases=(DataObjectBase,),
+                attrs={
+                    "__annotations__": {"files": List[str]},
+                    "__qualname__": f"{cls_name}.ListOfFiles",
+                },
+            ),
+        )
+        Directory = dataobject(package=package, name=f"{cls_name}Directory")(
+            _DataObjectBaseMetaClass.__new__(
+                _DataObjectBaseMetaClass,
+                name="Directory",
+                bases=(DataObjectBase,),
+                attrs={
+                    "__annotations__": {
+                        "dirname": str,
+                        "extension": str,
                     },
-                }
-            }
-        }
-        # TODO: Once full oneof support is implemented, this can move to using
-        #   the proper @dataclass style. In the meantime, the "schema" version
-        #   that uses JTD is still supported so that this oneof can be
-        #   dynamically declared. In order to make this work, the
-        #   _DataObjectBaseMetaClass needs to pre-know the right set of field
-        #   names for the oneof fields as if this were a dataclass. The values
-        #   of these annotations are unused since this will never be made into a
-        #   true @dataclass.
-        annotations = {
-            name.lower(): None
-            for name in schema["properties"]["data_stream"]["mapping"]
-        }
-        data_object = dataobject(schema=schema, package="caikit_data_model.runtime",)(
+                    "__qualname__": f"{cls_name}.Directory",
+                },
+            ),
+        )
+
+        # Create the outer class that encapsulates the Union (oneof) or the
+        # various types of input sources
+        data_object = dataobject(package=package)(
             _DataObjectBaseMetaClass.__new__(
                 _DataObjectBaseMetaClass,
                 name=cls_name,
                 bases=(DataObjectBase, DataStreamSourceBase),
                 attrs={
                     "ELEMENT_TYPE": data_element_type,
-                    "__annotations__": annotations,
+                    JsonData.__name__: JsonData,
+                    File.__name__: File,
+                    ListOfFiles.__name__: ListOfFiles,
+                    Directory.__name__: Directory,
+                    "__annotations__": {
+                        "data_stream": Union[
+                            Annotated[JsonData, OneofField(JsonData.__name__.lower())],
+                            Annotated[File, OneofField(File.__name__.lower())],
+                            Annotated[
+                                ListOfFiles, OneofField(ListOfFiles.__name__.lower())
+                            ],
+                            Annotated[
+                                Directory, OneofField(Directory.__name__.lower())
+                            ],
+                        ],
+                    },
                 },
             )
         )
+
+        # Add this data stream source to the common data model and the module
+        # where it was declared
         setattr(
             caikit.interfaces.common.data_model,
             cls_name,
@@ -288,10 +303,15 @@ def make_data_stream_source(data_element_type: Type) -> Type[DataBase]:
 
         # Add an init that sequences the initialization so that
         # DataStreamSourceBase is initialized after DataBase
-        orig_init = _DataObjectBaseMetaClass._make_init(data_object.fields)
+        orig_init = _make_oneof_init(data_object)
 
         def __init__(self, *args, **kwargs):
-            orig_init(self, *args, **kwargs)
+            try:
+                orig_init(self, *args, **kwargs)
+            except TypeError as err:
+                raise CaikitRuntimeException(
+                    grpc.StatusCode.INVALID_ARGUMENT, str(err)
+                ) from err
             DataStreamSourceBase.__init__(self)
 
         setattr(data_object, "__init__", __init__)


Test patch:
diff --git a/tests/core/data_model/data_backends/test_dict_backend.py b/tests/core/data_model/data_backends/test_dict_backend.py
index c62d554c..62b92727 100644
--- a/tests/core/data_model/data_backends/test_dict_backend.py
+++ b/tests/core/data_model/data_backends/test_dict_backend.py
@@ -128,6 +128,19 @@ def test_dict_backend_repeated_sub_message():
         assert sub_msg2.bar == data_dict["some_foos"][1]["bar"]
 
 
+def test_dict_backend_oneof():
+    """Make sure that a oneof can be correctly accessed from a backend"""
+    with temp_data_model(make_proto_def({"Foo": {"foo": "Union[str, int]"}})) as dm:
+        assert hasattr(dm, "Foo")
+        assert issubclass(dm.Foo, DataBase)
+
+        data_dict = {"foo": "asdf"}
+        backend = DictBackend(data_dict)
+        msg = dm.Foo.from_backend(backend)
+        assert msg.foo == "asdf"
+        assert msg.which_oneof("foo") == "foostr"
+
+
 def test_dict_backend_invalid_field_error():
     """Make sure that an AttributeError is raised if an invalid field is
     requested
diff --git a/tests/core/data_model/test_base.py b/tests/core/data_model/test_base.py
index 1a53218d..f684ad8e 100644
--- a/tests/core/data_model/test_base.py
+++ b/tests/core/data_model/test_base.py
@@ -24,6 +24,7 @@ import pytest
 # Local
 from caikit.core.data_model.base import DataBase
 from caikit.core.data_model.data_backends.base import DataModelBackendBase
+from caikit.core.data_model.dataobject import DataObjectBase
 from tests.data_model_helpers import (
     justify_script_string,
     make_proto_def,
@@ -130,6 +131,8 @@ def test_derived_class_no_import_side_effects():
                     class FakeDescriptor:
                         fields = []
                         fields_by_name = {}
+                        oneofs_by_name = {}
+                        oneofs = []
                         name = "Baz"
                         full_name = "foo.bar.Baz"
 
@@ -161,6 +164,86 @@ def test_derived_class_no_import_side_effects():
         assert issubclass(lib.Object, DataBase)
 
 
+def test_compiled_proto_init():
+    """Make sure that support for 'compiled' protos works cleanly without using
+    the dataobject wrapper
+    """
+    # pylint: disable=duplicate-code
+    with temp_data_model(
+        make_proto_def(
+            {
+                "ThingOne": {
+                    "foo": str,
+                    "bar": int,
+                }
+            },
+            mock_compiled=True,
+        )
+    ) as dm:
+        assert isinstance(dm.ThingOne, type)
+        assert issubclass(dm.ThingOne, DataBase)
+        assert not issubclass(dm.ThingOne, DataObjectBase)
+        assert dm.ThingOne.fields == ("foo", "bar")
+
+        # Test construction with positional args
+        inst = dm.ThingOne("foo", 1)
+        assert inst.foo == "foo"
+        assert inst.bar == 1
+
+        # Test construction with keyword args
+        inst = dm.ThingOne(foo="foo", bar=1)
+        assert inst.foo == "foo"
+        assert inst.bar == 1
+
+        # Test defaulting to None
+        inst = dm.ThingOne()
+        assert inst.foo is None
+        assert inst.bar is None
+
+        # Test error cases for construction
+        with pytest.raises(TypeError):
+            dm.ThingOne("foo", foo="bar")
+        with pytest.raises(TypeError):
+            dm.ThingOne(widget="qewr")
+        with pytest.raises(TypeError):
+            dm.ThingOne("foo", 1, 2)
+
+
+def test_compiled_proto_oneof():
+    """Make sure that support for 'compiled' protos works cleanly without using
+    the dataobject wrapper
+    """
+    # pylint: disable=duplicate-code
+    with temp_data_model(
+        make_proto_def(
+            {
+                "ThingOne": {
+                    "foo": "Union[str, int]",
+                }
+            },
+            mock_compiled=True,
+        )
+    ) as dm:
+        assert isinstance(dm.ThingOne, type)
+        assert issubclass(dm.ThingOne, DataBase)
+        assert not issubclass(dm.ThingOne, DataObjectBase)
+        assert set(dm.ThingOne.fields) == {"foostr", "fooint"}
+
+        # Construct with the oneof name
+        inst = dm.ThingOne(foo=1)
+        assert inst.foo == 1
+        assert inst.which_oneof("foo") == "fooint"
+
+        # Construct with field name
+        inst = dm.ThingOne(foostr="asdf")
+        assert inst.foo == "asdf"
+        assert inst.which_oneof("foo") == "foostr"
+
+        # Conflicting args
+        with pytest.raises(TypeError):
+            dm.ThingOne(foo=1, fooster="asdf")
+
+
 ##################
 ## Data Backend ##
 ##################
@@ -224,6 +307,26 @@ def test_uncached_backend():
         assert backend.access_count("bar") == 3
 
 
+def test_invalid_attribute_no_backend():
+    """Make sure that when created without a backend and without proper
+    initialization, an AttributeError is raised
+    """
+    # pylint: disable=duplicate-code
+    with temp_data_model(
+        make_proto_def(
+            {
+                "ThingOne": {
+                    "foo": str,
+                    "bar": int,
+                }
+            }
+        )
+    ) as dm:
+        msg = dm.ThingOne.__new__(dm.ThingOne)
+        with pytest.raises(AttributeError):
+            msg.foo
+
+
 ##################
 ## Default Init ##
 ##################
diff --git a/tests/core/data_model/test_dataobject.py b/tests/core/data_model/test_dataobject.py
index 2ca6e75d..dd1fd416 100644
--- a/tests/core/data_model/test_dataobject.py
+++ b/tests/core/data_model/test_dataobject.py
@@ -113,22 +113,6 @@ def test_dataobject_native_types():
     assert inst.bar is None
 
 
-def test_dataobject_jtd():
-    """Make sure that a simple usage of dataobject using a full JTD schema works
-    as expected
-    """
-
-    @dataobject(schema={"properties": {"foo": {"type": "string"}}})
-    class Foo(DataObjectBase):
-        foo: str
-
-    assert check_field_type(Foo.get_proto_class(), "foo", "TYPE_STRING")
-    inst = Foo(foo="test")
-    assert inst.foo == "test"
-    inst = Foo()
-    assert inst.foo is None
-
-
 def test_dataobject_nested_objects():
     """Make sure that nested objects are handled correctly"""
 
@@ -279,15 +263,6 @@ def test_dataobject_obj_refs_with_optional_types():
         )
 
 
-def test_dataobject_invalid_schema():
-    """Make sure that a ValueError is raised on an invalid schema"""
-    with pytest.raises(ValueError):
-        # pylint: disable=unused-variable
-        @dataobject(schema="Foo")
-        class Foo:
-            pass
-
-
 def test_dataobject_additional_methods():
     """Make sure that additional methods on wrapped classes (for messages and
     enums) are preserved
@@ -372,75 +347,11 @@ def test_render_dataobject_protos_no_dir():
         }
 
 
-def test_dataobject_with_discriminator():
-    """Make sure that adding a discriminator works as expected"""
-
-    @dataobject(
-        schema={
-            "properties": {
-                "data_stream": {
-                    "discriminator": "data_reference_type",
-                    "mapping": {
-                        "Foo": {
-                            "properties": {
-                                "data": {
-                                    "elements": {"type": "string"},
-                                },
-                            },
-                        },
-                        "Bar": {"properties": {"data": {"type": "string"}}},
-                        "Baz": {
-                            "properties": {
-                                "data": {
-                                    "elements": {"type": "string"},
-                                },
-                            },
-                        },
-                        "Bat": {
-                            "properties": {
-                                "data1": {"type": "string"},
-                                "data2": {"type": "string"},
-                            }
-                        },
-                    },
-                }
-            }
-        }
-    )
-    class BazObj(DataObjectBase):
-        foo: str
-        bar: str
-        baz: str
-        bat: str
-
-    # proto tests
-    foo1 = BazObj(foo=BazObj.Foo(data=["hello"]))
-    proto_repr_foo = foo1.to_proto()
-    assert proto_repr_foo.foo.data == ["hello"]
-    assert BazObj.from_proto(proto=proto_repr_foo).to_proto() == proto_repr_foo
-
-    bar1 = BazObj(foo=BazObj.Foo(data=["hello"]), bar=BazObj.Bar(data="world"))
-    proto_repr_bar = bar1.to_proto()
-    assert proto_repr_bar.bar.data == "world"
-
-    # json tests
-    foo1 = BazObj(foo=BazObj.Foo(data=["hello"]))
-    json_repr_foo = foo1.to_json()
-    assert json.loads(json_repr_foo) == {
-        "foo": {"data": ["hello"]},
-        "bar": None,
-        "baz": None,
-        "bat": None,
-    }
-
-
 def test_dataobject_with_oneof():
     """Make sure that using a Union to create a oneof works as expected"""
 
     @dataobject
     class BazObj(DataObjectBase):
-        _private_slots = ("_which_oneof_datastream",)
-
         @dataobject
         class Foo(DataObjectBase):
             data: List[str]
@@ -454,47 +365,96 @@ def test_dataobject_with_oneof():
             Annotated[Bar, FieldNumber(2), OneofField("bar")],
         ]
 
-        def __getattr__(self, name):
-            if name == "data_stream":
-                if self._which_oneof_datastream == "foo":
-                    return self.foo
-                elif self._which_oneof_datastream == "bar":
-                    return self.bar
-                return None
-            if name == "_foo":
-                if self._which_oneof_datastream == "foo":
-                    return self._data_stream
-            if name == "_bar":
-                if self._which_oneof_datastream == "bar":
-                    return self._data_stream
-
-        def __init__(self, *args, **kwargs):
-            if "foo" in kwargs:
-                self._which_oneof_datastream = "foo"
-                self._data_stream = kwargs["foo"]
-            if "bar" in kwargs:
-                self._which_oneof_datastream = "bar"
-                self._data_stream = kwargs["bar"]
-
-    # proto tests
+    # Construct with oneof field name
     foo1 = BazObj(foo=BazObj.Foo(data=["hello"]))
     assert isinstance(foo1.data_stream, BazObj.Foo)
+    assert foo1.which_oneof("data_stream") == "foo"
+    assert foo1.foo is foo1.data_stream
+    assert foo1.bar is None
+
+    # Test other oneof field name
+    bar1 = BazObj(bar=BazObj.Bar(data="world"))
+    assert isinstance(bar1.data_stream, BazObj.Bar)
+    assert bar1.which_oneof("data_stream") == "bar"
+    assert bar1.bar is bar1.data_stream
+    assert bar1.foo is None
+
+    # Test proto round trip
     proto_repr_foo = foo1.to_proto()
     assert proto_repr_foo.foo.data == ["hello"]
     assert BazObj.from_proto(proto=proto_repr_foo).to_proto() == proto_repr_foo
-
-    bar1 = BazObj(foo=BazObj.Foo(data=["hello"]), bar=BazObj.Bar(data="world"))
-    assert isinstance(bar1.data_stream, BazObj.Bar)
     proto_repr_bar = bar1.to_proto()
     assert proto_repr_bar.bar.data == "world"
 
-    # json tests
-    foo1 = BazObj(foo=BazObj.Foo(data=["hello"]))
+    # Test json round trip
     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
+
+    # Test setattr
+    foo1.bar = BazObj.Bar(data="it's a bar")
+    assert foo1.which_oneof("data_stream") == "bar"
+    assert foo1.data_stream is foo1.bar
+    assert foo1.foo is None
+
+    # Construct with oneof name
+    foo2 = BazObj(data_stream=BazObj.Foo(data=["some", "foo"]))
+    assert foo2.data_stream.data == ["some", "foo"]
+    assert foo2.bar is None
+    assert foo2.foo is foo2.data_stream
+    assert foo2.which_oneof("data_stream") == "foo"
+
+    # Assign with oneof name
+    foo2.data_stream = BazObj.Bar(data="asdf")
+    assert foo2.foo is None
+    assert foo2.bar is foo2.data_stream
+    assert foo2.which_oneof("data_stream") == "bar"
+
+    # Construct with positional oneof name
+    foo2 = BazObj(BazObj.Foo(data=["some", "foo"]))
+    assert foo2.data_stream.data == ["some", "foo"]
+    assert foo2.bar is None
+    assert foo2.foo is foo2.data_stream
+    assert foo2.which_oneof("data_stream") == "foo"
+
+    foo3 = BazObj()
+    assert foo3.foo is None
+    assert foo3.bar is None
+    assert foo3.data_stream is None
+    assert foo3.which_oneof("data_stream") == None
+    # Invalid constructors
+    with pytest.raises(TypeError):
+        BazObj(BazObj.Foo(), foo=BazObj.Foo())
+    with pytest.raises(TypeError):
+        BazObj(data_stream=BazObj.Foo(), foo=BazObj.Foo())
+    with pytest.raises(TypeError):
+        BazObj(foo=BazObj.Foo(), bar=BazObj.Bar())
+
+
+def test_dataobject_with_same_type_of_oneof():
+    """Make sure that using a Union to create a oneof with the same types works as expected"""
+
+    @dataobject
+    class Foo(DataObjectBase):
+        foo: Union[
+            Annotated[bool, FieldNumber(10), OneofField("foo_bool1")],
+            Annotated[bool, FieldNumber(20), OneofField("foo_bool2")],
+        ]
+
+    # if the fields are of the same type, then by default the first one is set
+    foo1 = Foo(True)
+    assert foo1.which_oneof("foo") == "foo_bool1"
+    assert foo1.foo_bool1
+    assert foo1.foo_bool2 == None
+
+    # unless set explicitly
+    foo2 = Foo(foo_bool2=True)
+    assert foo2.which_oneof("foo") == "foo_bool2"
+    assert foo2.foo_bool1 == None
+    assert foo2.foo_bool2
 
 
 def test_dataobject_round_trip_json():
diff --git a/tests/data_model_helpers.py b/tests/data_model_helpers.py
index 51f11d08..81e7de83 100644
--- a/tests/data_model_helpers.py
+++ b/tests/data_model_helpers.py
@@ -147,24 +147,61 @@ def _get_proto_val_name(field_val) -> str:
     raise RuntimeError(f"Invalid field type specifier: {field_val}")
 
 
-def make_proto_def(message_specs: Dict[str, dict], pkg_suffix: str = None) -> str:
+def make_proto_def(
+    message_specs: Dict[str, dict],
+    pkg_suffix: str = None,
+    mock_compiled: bool = False,
+) -> str:
     """Helper for writing a syntatically correct protobufs file"""
     if pkg_suffix is None:
         pkg_suffix = _random_package_suffix()
     package_name = f"{caikit.core.data_model.CAIKIT_DATA_MODEL}.{pkg_suffix}"
     out = justify_script_string(
         """
-        from typing import Dict, List
-        from caikit.core.data_model import DataObjectBase, dataobject
-
+        from typing import Dict, List, Union
         """
     )
-    for message_name, message_spec in message_specs.items():
-        msg_str = (
-            f'\n@dataobject("{package_name}")\nclass {message_name}(DataObjectBase):\n'
+    if mock_compiled:
+        out += justify_script_string(
+            """
+            from caikit.core.data_model import DataBase
+            from py_to_proto import dataclass_to_proto, descriptor_to_message_class
+            from dataclasses import dataclass
+            """
         )
+    else:
+        out += "\nfrom caikit.core.data_model import DataObjectBase, dataobject\n"
+    for message_name, message_spec in message_specs.items():
+        type_annotations = ""
         for field_name, field_type in message_spec.items():
-            msg_str += f"    {field_name}: {_get_proto_val_name(field_type)}\n"
-        msg_str += "\n"
+            type_annotations += f"    {field_name}: {_get_proto_val_name(field_type)}\n"
+
+        # Manually create the proto class and use the precompiled proto style
+        if mock_compiled:
+            dataclass_name = f"_{message_name}"
+            proto_name = f"{dataclass_name}_proto"
+            msg_str = f"\n@dataclass\nclass {dataclass_name}:\n"
+            msg_str += type_annotations
+            msg_str += justify_script_string(
+                f"""
+
+                {proto_name} = descriptor_to_message_class(
+                    dataclass_to_proto(
+                        dataclass_={dataclass_name},
+                        package="{package_name}",
+                        name="{message_name}",
+                    )
+                )
+
+                class {message_name}(DataBase):
+                    _proto_class = {proto_name}
+                """
+            )
+
+        # Add as a dataobject
+        else:
+            msg_str = f'\n@dataobject("{package_name}")\nclass {message_name}(DataObjectBase):\n'
+            msg_str += type_annotations
+
         out += msg_str
     return out