caikit__caikit-233

The user is reporting a bug in `caikit` where inheriting from a `DataObjectBase` class and overriding methods without adding new fields causes a `TypeError: All DataObjectBase classes must follow dataclass syntax`.

Looking at the sample code provided:
```py
from caikit.core import DataObjectBase, dataobject

@dataobject
class Base(DataObjectBase):
    foo: int
    def doit(self):
        return self.foo * 2

@dataobject
class Derived(Base):
    def doit(self):
        return self.foo * 3
```
This causes a `TypeError` during the definition of `Derived`.

The error comes from `/Users/ghart/Projects/github/caikit/caikit2/caikit/core/data_model/dataobject.py` line 88 (path in the issue report, likely slightly different in the current workspace).

I need to find this code and understand why it's raising the `TypeError`.

Phase 1: Pinpoint Code & Tests.
I'll search for `DataObjectBase` and `dataobject` in `caikit/core/data_model/dataobject.py`.

{
  "command": "grep -rn --include=\"*.py\" \"DataObjectBase\" caikit/core/data_model/dataobject.py"
}

{
  "status": "ok",
  "stdout": "72:class _DataObjectBaseMetaClass(_DataBaseMetaClass):\n85:        if name != \"DataObjectBase\":\n89:                    \"All DataObjectBase classes must follow dataclass syntax\"\n97:class DataObjectBase(DataBase, metaclass=_DataObjectBaseMetaClass):\n112:    class MyDataObject(DataObjectBase):\n135:            issubclass(cls, (DataObjectBase, Enum)),\n136:            \"{} must inherit from DataObjectBase/Enum when using @dataobject\",\n295:    if issubclass(cls, DataObjectBase):\n",
  "stderr": "",
  "exit_code": 0
}

Okay, the error is likely in `caikit/core/data_model/dataobject.py` within `_DataObjectBaseMetaClass.__new__`.

Let's read the file.

{
  "end_line": 120,
  "filepath": "caikit/core/data_model/dataobject.py",
  "start_line": 70
}

{
  "status": "ok",
  "filepath": "caikit/core/data_model/dataobject.py",
  "content": "\n\nclass _DataObjectBaseMetaClass(_DataBaseMetaClass):\n    \"\"\"This metaclass is used for the DataObject base class so that all data\n    objects can delay the creation of their proto class until after the\n    metaclass has been instantiated.\n    \"\"\"\n\n    def __new__(mcs, name, bases, attrs):\n        \"\"\"When instantiating a new DataObject class, the proto class will not\n        yet have been generated, but the set of fields will be known since the\n        class will be the raw input representation of a @dataclass\n        \"\"\"\n\n        # Get the annotations that will go into the dataclass\n        if name != \"DataObjectBase\":\n            field_names = attrs.get(\"__annotations__\")\n            if field_names is None:\n                raise TypeError(\n                    \"All DataObjectBase classes must follow dataclass syntax\"\n                )\n            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names\n\n        # Delegate to the base metaclass\n        return super().__new__(mcs, name, bases, attrs)\n\n\nclass DataObjectBase(DataBase, metaclass=_DataObjectBaseMetaClass):\n    \"\"\"A DataObject is a data model class that is backed by a @dataclass.\n\n    Data model classes that use the @dataobject decorator must derive from this\n    base class.\n    \"\"\"\n\n\ndef dataobject(*args, **kwargs) -> Callable[[Type], Type[DataBase]]:\n    \"\"\"The @dataobject decorator can be used to define a Data Model object's\n    schema inline with the definition of the python class rather than needing to\n    bind to a pre-compiled protobufs class. For example:\n\n    @dataobject(\"foo.bar\")\n    @dataclass\n    class MyDataObject(DataObjectBase):\n        '''My Custom Data Object'''\n        foo: str\n        bar: int\n\n    NOTE: The wrapped class must NOT inherit directly from DataBase. That\n        inheritance will be added by this decorator, but if it is written\n        directly, the metaclass that links protobufs to the class will be called\n        before this decorator can auto-gen the protobufs class.",
  "start_line": 70,
  "end_line": 120,
  "total_lines": 386,
  "is_truncated": false
}

The code in `_DataObjectBaseMetaClass.__new__` is:

```python
    def __new__(mcs, name, bases, attrs):
        """When instantiating a new DataObject class, the proto class will not
        yet have been generated, but the set of fields will be known since the
        class will be the raw input representation of a @dataclass
        """

        # Get the annotations that will go into the dataclass
        if name != "DataObjectBase":
            field_names = attrs.get("__annotations__")
            if field_names is None:
                raise TypeError(
                    "All DataObjectBase classes must follow dataclass syntax"
                )
            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names

        # Delegate to the base metaclass
        return super().__new__(mcs, name, bases, attrs)
```

If I create a class `Derived(Base)` where `Base` is a `DataObjectBase` and `Derived` doesn't define any new fields (no annotations), `attrs.get("__annotations__")` will be `None`, and it will raise the `TypeError`.

However, `Derived` still has the fields from `Base` because it inherits from it. The design seems to expect `__annotations__` to be present on every `DataObjectBase` subclass.

If we want to support functional inheritance where a subclass only overrides methods but not fields, we should probably check if any of the base classes also have annotations if the current class doesn't. Wait, `attrs` only contains the attributes of the class being created, not its bases.

If `Derived` inherits from `Base`, `Base` already has the annotations. But `_DataObjectBaseMetaClass.__new__` is called for `Derived` as well.

The problem is that `attrs.get("__annotations__")` is `None` for `Derived`.

If `Derived` is intended to be a `DataObjectBase` and inherits from another `DataObjectBase`, it should be allowed to not have its own annotations, provided one of its bases does (or maybe it's just always allowed if it inherits from `DataObjectBase`?).

Wait, if it inherits from `Base`, and `Base` is a `DataObjectBase`, then `Derived` is also a `DataObjectBase`.

The line `if name != "DataObjectBase":` is meant to avoid this check for the base class itself.

If `Derived` inherits from `Base`, we might want to check if `Base` has animations. But `bases` (the second argument to `__new__`) contains the base classes.

However, `Base` might have had its annotations stored in `_DataBaseMetaClass._FWD_DECL_FIELDS` (which is what the code does).

Let's see what `_DataBaseMetaClass._FWD_DECL_FIELDS` is. It's likely a string key used in the attrs dictionary of the class.

Wait, if `Derived` inherits from `Base`, it doesn't necessarily need its own `__annotations__` if it doesn't add new fields.

The fix should be: if `__annotations__` is missing, check if any of the bases are `DataObjectBase` or have `_FWD_DECL_FIELDS`. No, that's not quite right. If any base is a `DataObjectBase` (and thus went through this metaclass), it probably already defined the fields.

Wait, if `Derived` inherits from `Base`, then `Base` had annotations. If `Derived` doesn't have annotations, it just uses `Base`'s annotations.

Actually, if `Derived` doesn't have `__annotations__`, it's just not adding new fields.

Let's try to reproduce this first.

{
  "command": "python3 - <<'PY'\nfrom caikit.core import DataObjectBase, dataobject\nfrom typing import Any\n\n@dataobject\nclass Base(DataObjectBase):\n    foo: int\n    def doit(self):\n        return self.foo * 2\n\ntry:\n    @dataobject\n    class Derived(Base):\n        def doit(self):\n            return self.foo * 3\n    print(\"Success\")\nexcept TypeError as e:\n    print(f\"Caught expected error: {e}\")\nexcept Exception as e:\n    print(f\"Caught unexpected error: {type(e).__name__}: {e}\")\nPY\n"
}

{
  "status": "ok",
  "stdout": "Caught expected error: All DataObjectBase classes must follow dataclass syntax\n",
  "stderr": "<function register_backend_type at 0x740a3cf66290> is still in the BETA phase and subject to change!\n",
  "exit_code": 0
}

Reproduced.

Now, how to fix it.

The `_DataObjectBaseMetaClass.__new__` method should allow classes that inherit from `DataObjectBase` to not have their own `__annotations__` if they are inheriting from another class that does.

But wait, if they inherit from `DataObjectBase`, they should probably just be allowed to not have annotations if at least one base class is a `DataObjectBase`.

Wait, `DataObjectBase` itself does not have `__annotations__`.

If `Base` inherits from `DataObjectBase`, `Base` must have `__annotations__`.
If `Derived` inherits from `Base`, `Derived` can have `__annotations__` (if it adds fields) or not.

So the condition `if field_names is None: raise TypeError(...)` is too strict for subclasses of other `DataObjectBase` classes.

Let's check the bases.

```python
    def __new__(mcs, name, bases, attrs):
        # ...
        if name != "DataObjectBase":
            field_names = attrs.get("__annotations__")
            if field_names is None:
                # Check if any base is a DataObjectBase
                # But DataObjectBase is just a class.
                # We can check if any base is a subclass of DataObjectBase.
                # However, DataObjectBase is not yet fully defined when _DataObjectBaseMetaClass.__new__ 
                # is called for DataObjectBase itself. 
                # But here we are in Derived(Base) where Base is already defined.
                
                # Let's see...
                pass
```

Wait, `DataObjectBase` is defined *after* `_DataObjectBaseMetaClass`.

When `Base` is created:
`name="Base"`, `bases=(DataObjectBase,)`, `attrs={...}`.
`DataObjectBase` is already defined.

When `Derived` is created:
`name="Derived"`, `bases=(Base,)`, `attrs={...}`.

If we check if any base is a subclass of `DataObjectBase`, it should work for `Derived`.

Wait, what about `Base`? `Base` inherits from `DataObjectBase`. Is `DataObjectBase` a subclass of `DataObjectBase`? No, usually `issubclass(C, C)` is True.

So for `Base`: `issubclass(DataObjectBase, DataObjectBase)` is True.
But for `Base`, `field_names = attrs.get("__annotations__")` should be present. If it's not, then `Base` (which is the first level of data object) MUST define the fields.

Wait, if I do:
```python
@dataobject
class Base(DataObjectBase):
    pass
```
This should probably also fail because it doesn't define any fields.

The issue report says: "When creating a class that inherits from a DataObjectBase and overrides parent functions without adding or changing any of the schema, an exception is raised".

So if `Derived` inherits from `Base`, and `Base` already defined some schema, `Derived` should be allowed to not define any new schema.

So the logic should be:
If `__annotations__` is missing:
  If any base is a subclass of `DataObjectBase` (and it's not `DataObjectBase` itself), then it's okay.
  Actually, if any base is a subclass of `DataObjectBase`, it means we are inheriting from a data object, so we have a schema.

Wait, what if I have:
```python
class Base(DataObjectBase): # without @dataobject
    foo: int
```
Wait, the `@dataobject` decorator is also used.

Let's look at `@dataobject` implementation.

{
  "end_line": 320,
  "filepath": "caikit/core/data_model/dataobject.py",
  "start_line": 120
}

{
  "status": "ok",
  "filepath": "caikit/core/data_model/dataobject.py",
  "content": "        before this decorator can auto-gen the protobufs class.\n\n    Args:\n        package:  str\n            The package name to use for the generated protobufs class\n\n    Returns:\n        decorator:  Callable[[Type], Type[DataBase]]\n            The decorator function that will wrap the given class\n    \"\"\"\n\n    def decorator(cls: Type) -> Type[DataBase]:\n        # Make sure that the wrapped class does NOT inherit from DataBase\n        error.value_check(\n            \"<COR95184230E>\",\n            issubclass(cls, (DataObjectBase, Enum)),\n            \"{} must inherit from DataObjectBase/Enum when using @dataobject\",\n            cls.__name__,\n        )\n\n        # Add the package to the kwargs\n        kwargs.setdefault(\"package\", package)\n\n        # If it's not an enum, fill in any missing field defaults as None\n        # and make sure it's a dataclass\n        if not issubclass(cls, Enum):\n            log.debug2(\"Wrapping data class %s\", cls)\n            user_defined_defaults = {}\n            for annotation in getattr(cls, \"__annotations__\", {}):\n                user_defined_default = getattr(cls, annotation, dataclasses.MISSING)\n                if user_defined_default == dataclasses.MISSING:\n                    log.debug3(\"Filling in None default for %s.%s\", cls, annotation)\n                    setattr(cls, annotation, None)\n                else:\n                    user_defined_defaults[annotation] = user_defined_default\n            # If the current __init__ is auto-generated by dataclass, remove\n            # it so that a new one is created with the new defaults\n            if _has_dataclass_init(cls):\n                log.debug3(\"Resetting default dataclass init\")\n                delattr(cls, \"__init__\")\n            cls = dataclasses.dataclass(cls)\n            setattr(cls, _USER_DEFINED_DEFAULTS, user_defined_defaults)\n\n        descriptor = _dataobject_to_proto(dataclass_=cls, **kwargs)\n\n        # Create the message class from the dataclass\n        proto_class = py_to_proto.descriptor_to_message_class(descriptor)\n        _AUTO_GEN_PROTO_CLASSES.append(proto_class)\n\n        # Add enums to the global enums module\n        for enum_class in _get_all_enums(proto_class):\n            log.debug2(\"Importing enum [%s]\", enum_class.DESCRIPTOR.name)\n            enums.import_enum(enum_class)\n\n        # Declare the merged class that binds DataBase to the wrapped class with\n        # this generated proto class\n        if isinstance(proto_class, type):\n            setattr(cls, \"_proto_class\", proto_class)\n            cls = _make_data_model_class(proto_class, cls)\n\n            # If this was a default-generated dataclass __init__ and there are\n            # any oneofs, we need to augment the __init__ to support kwargs for\n            # the individual fields\n            if _has_dataclass_init(cls) and cls._fields_oneofs_map:\n                setattr(cls, \"__init__\", _make_oneof_init(cls))\n\n        else:\n            enums.import_enum(proto_class, cls)\n            setattr(cls, \"_proto_enum\", proto_class)\n\n        # Return the decorated class\n        return cls\n\n    # If called without the function invocation, fill in the default argument\n    if args and callable(args[0]):\n        assert not kwargs, \"This shouldn't happen!\"\n        package = CAIKIT_DATA_MODEL\n        return decorator(args[0])\n\n    # Pull the package as an arg or a keyword arg\n    if args:\n        package = args[0]\n        if \"package\" in kwargs:\n            raise TypeError(\"Got multiple values for argument 'package'\")\n    else:\n        package = kwargs.get(\"package\", CAIKIT_DATA_MODEL)\n    return decorator\n\n\ndef render_dataobject_protos(interfaces_dir: str):\n    \"\"\"Write out protobufs files for all proto classes generated from dataobjects\n    to the target interfaces directory\n\n    Args:\n        interfaces_dir:  str\n            The target directory (must already exist)\n    \"\"\"\n    for proto_class in _AUTO_GEN_PROTO_CLASSES:\n        proto_class.write_proto_file(interfaces_dir)\n\n\n## Implementation Details ######################################################\n\n\ndef _dataobject_to_proto(*args, **kwargs):\n    kwargs.setdefault(\"type_mapping\", DATAOBJECT_PY_TO_PROTO_TYPES)\n    return _DataobjectConverter(*args, **kwargs).descriptor\n\n\nclass _DataobjectConverter(DataclassConverter):\n    \"\"\"Augment the dataclass converter to be able to pull descriptors from\n    existing data objects\n    \"\"\"\n\n    def get_concrete_type(self, entry: Any) -> Any:\n        \"\"\"Also include data model classes and enums as concrete types\"\"\"\n        unwrapped = self._resolve_wrapped_type(entry)\n        if (\n            isinstance(unwrapped, type)\n            and issubclass(unwrapped, DataBase)\n            and unwrapped._proto_class is not None\n        ) or hasattr(unwrapped, \"_proto_enum\"):\n            return entry\n        return super().get_concrete_type(entry)\n\n    def get_descriptor(self, entry: Any) -> Any:\n        \"\"\"Unpack data model classes and enums to their descriptors\"\"\"\n        entry = self._resolve_wrapped_type(entry)\n        if isinstance(entry, type) and issubclass(entry, DataBase):\n            return entry._proto_class.DESCRIPTOR\n        proto_enum = getattr(entry, \"_proto_enum\", None)\n        if proto_enum is not None:\n            return proto_enum.DESCRIPTOR\n        return super().get_descriptor(entry)\n\n    def get_optional_field_names(self, entry: Any) -> List[str]:\n        \"\"\"Get the names of any fields which are optional. This will be any\n        field that has a user-defined default or is marked as Optional[]\n        \"\"\"\n        optional_fields = list(getattr(entry, _USER_DEFINED_DEFAULTS, {}))\n        for field_name, field in entry.__dataclass_fields__.items():\n            if (\n                field_name not in optional_fields\n                and self._is_python_optional(field.type) is not None\n            ):\n                optional_fields.append(field_name)\n        return optional_fields\n\n    @staticmethod\n    def _is_python_optional(entry: Any) -> Any:",
  "start_line": 120,
  "end_line": 269,
  "total_lines": 386,
  "is_truncated": true
}

The `@dataobject` decorator calls `dataclasses.dataclass(cls)`.

The `_DataObjectBaseMetaClass` is used when the class is being created.

```python
class _DataObjectBaseMetaClass(_DataBaseMetaClass):
    # ...
    def __new__(mcs, name, bases, attrs):
        # ...
        if name != "DataObjectBase":
            field_names = attrs.get("__annotations__")
            if field_names is None:
                raise TypeError(
                    "All DataObjectBase classes must follow dataclass syntax"
                )
            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names

        # Delegate to the base metaclass
        return super().__new__(mcs, name, bases, attrs)
```

If `Derived` inherits from `Base`, and `Base` is already a `DataObjectBase`, then `Derived` doesn't need its own `__annotations__` if it doesn't add new fields.

However, if we just remove the `TypeError`, we need to make sure `attrs[_DataBaseMetaClass._FWD_DECL_FIELDS]` is still set correctly.

If `Derived` has no annotations, it should probably inherit the fields from `Base`.

Wait, `_DataBaseMetaClass._FWD_DECL_FIELDS` is used to store the field names.

If `Derived` has no annotations, `attrs.get("__annotations__")` is `None`.

Should we check if any base class has annotations or `_FWD_DECL_FIELDS`?

If `Derived` inherits from `Base`, `Base` already has the annotations in its own `__annotations__` (if it's a dataclass) and it might have set `_FWD_DECL_FIELDS` on itself.

Wait, `_DataObjectBaseMetaClass.__new__` sets `attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names`. This is setting it on the *class* being created.

So `Base._FWD_DECL_FIELDS` (where `_FWD_DECL_FIELDS` is the value of `_DataBaseMetaClass._FWD_DECL_FIELDS`) will be set.

Let's see what `_DataBaseMetaClass._FWD_DECL_FIELDS` is. I'll search for it.

{
  "command": "grep -rn \"_FWD_DECL_FIELDS\" caikit/core/data_model/dataobject.py"
}

{
  "status": "ok",
  "stdout": "91:            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names\n",
  "stderr": "",
  "exit_code": 0
}

It's only used there. Let's see where `_DataBaseMetaClass` is defined. It's not in `caikit/core/data_model/dataobject.py`. It must be in a base class.

Wait, `_DataObjectBaseMetaClass` inherits from `_DataBaseMetaClass`.

Let's search for `class _DataBaseMetaClass` in the repository.

{
  "command": "grep -rn \"class _DataBaseMetaClass\" caikit/core/data_model/"
}

{
  "status": "ok",
  "stdout": "caikit/core/data_model/base.py:46:class _DataBaseMetaClass(type):\n",
  "stderr": "",
  "exit_code": 0
}

Let's read `caikit/core/data_model/base.py`.

{
  "filepath": "caikit/core/data_model/base.py"
}

{
  "status": "ok",
  "filepath": "caikit/core/data_model/base.py",
  "content": "# Copyright The Caikit Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n\"\"\"Base classes and functionality for all data structures.\n\"\"\"\n\n# metaclass-generated field members cannot be detected by pylint\n# pylint: disable=no-member\n\n# Standard\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom typing import Any, Dict, Optional, Type, Union\nimport base64\nimport json\n\n# Third Party\nfrom google.protobuf import json_format\nfrom google.protobuf.descriptor import Descriptor\nfrom google.protobuf.internal import type_checkers as proto_type_checkers\nfrom google.protobuf.message import Message as ProtoMessageType\n\n# First Party\nimport alog\n\n# Local\nfrom ..toolkit.errors import error_handler\nfrom . import enums, json_dict\n\nlog = alog.use_channel(\"DATAM\")\nerror = error_handler.get(log)\n\n\nclass _DataBaseMetaClass(type):\n    \"\"\"Meta class for all structures in the data model.\"\"\"\n\n    # store a registry of all classes that use this metaclass, i.e.,\n    # all classes that extend DataBase. This is used for constructing new\n    # instances by name without having to introspect all modules in data_model.\n    class_registry = {}\n\n    # This sentinel value is used to determine whether a given attribute is\n    # present on a class without doing `getattr` twice in the case where the\n    # attribute does exist.\n    _MISSING_ATTRIBUTE = \"missing attribute\"\n\n    # Special attribute used to communicate that the proto fields are forward\n    # declared and will be populated after the metaclass has completed\n    # construction.\n    _FWD_DECL_FIELDS = \"__fwd_decl_fields__\"\n\n    # Special instance attributes that an instance of a class derived from\n    # DataBase may have. These are added to __slots__.\n    _BACKEND_ATTR = \"_backend\"\n    _WHICH_ONEOF_ATTR = \"_which_oneof\"\n\n    def __new__(mcs, name, bases, attrs):\n        \"\"\"When constructing a new data model class, we set the 'fields' class variable from the\n        protobufs descriptor and then set the '__slots__' magic class attribute to fields.  This\n        provides two benefits: (a) performance is improved since the classes only need to know\n        about these attributes (b) it helps to enforce that all member variables in these classes\n        are described in the protobufs.\n\n        Note:  If you want to add a variable for internal use that is not described in the\n        protobufs, it can be named in the tuple class variable _private_slots and will\n        automatically be added to __slots__.\n        \"\"\"\n\n        # Protobufs fields can be divided into these categories, which are used\n        # to automatically determine appropriate behavior in a number of methods\n        attrs[\"full_name\"] = name\n        attrs[\"fields_enum_map\"] = {}\n        attrs[\"fields_enum_rev\"] = {}\n        attrs[\"_fields_oneofs_map\"] = {}\n        attrs[\"_fields_to_oneof\"] = {}\n        attrs[\"_fields_map\"] = ()\n        attrs[\"_fields_message\"] = ()\n        attrs[\"_fields_message_repeated\"] = ()\n        attrs[\"_fields_enum\"] = ()\n        attrs[\"_fields_enum_repeated\"] = ()\n        attrs[\"_fields_primitive\"] = ()\n        attrs[\"_fields_primitive_repeated\"] = ()\n\n        # Look for the set of fields either from a predefined protobuf class or\n        # from a forward declaration from @dataobject\n        fields = ()\n        proto_class = None\n        if name not in [\"DataBase\", \"DataObjectBase\"]:\n            # Look for a precompiled proto class and if found, parse its\n            # descriptor\n            proto_class = attrs.get(\"_proto_class\")\n            if proto_class is not None:\n                all_oneof_fields = [\n                    field.name\n                    for oneof in proto_class.DESCRIPTOR.oneofs\n                    for field in oneof.fields\n                ]\n                fields = tuple(\n                    (\n                        field\n                        for field in proto_class.DESCRIPTOR.fields_by_name\n                        if field not in all_oneof_fields\n                    )\n                ) + tuple(proto_class.DESCRIPTOR.oneofs_by_name)\n\n            # Otherwise, we need to get the fields from a \"special\" attribute\n            else:\n                fields = attrs.pop(mcs._FWD_DECL_FIELDS, None)\n                log.debug4(\n                    \"Using dataclass forward declaration fields %s for %s\", fields, name\n                )\n                error.value_check(\n                    \"<COR49310991E>\",\n                    fields is not None,\n                    \"No proto class found for {}\",\n                    name,\n                )\n        attrs[\"fields\"] = fields\n        attrs[\"_proto_class\"] = proto_class\n\n        # Look if any private slots are declared as class variables\n        private_slots = attrs.setdefault(\"_private_slots\", ())\n\n        # Class slots are fields + private slots, this prevents other\n        # member attributes from being set and also improves performance\n        attrs[\"__slots__\"] = tuple(\n            [f\"_{field}\" for field in fields]\n            + list(private_slots)\n            + [mcs._BACKEND_ATTR, mcs._WHICH_ONEOF_ATTR]\n        )\n\n        # Create the instance of the type\n        instance = super().__new__(mcs, name, bases, attrs)\n\n        # If there's a valid proto class, perform proto descriptor parsing\n        if proto_class is not None:\n            mcs.parse_proto_descriptor(instance)\n",
  "start_line": 1,
  "end_line": 150,
  "total_lines": 990,
  "is_truncated": true
}

In `_DataBaseMetaClass.__new__`:

```python
        # Look for the set of fields either from a predefined protobuf class or
        # from a forward declaration from @dataobject
        fields = ()
        proto_class = None
        if name not in ["DataBase", "DataObjectBase"]:
            # ...
            # Otherwise, we need to get the fields from a "special" attribute
            else:
                fields = attrs.pop(mcs._FWD_DECL_FIELDS, None)
                log.debug4(
                    "Using dataclass forward declaration fields %s for %s", fields, name
                )
                error.value_check(
                    "<COR49310991E>",
                    fields is not None,
                    "No proto class found for {}",
                    name,
                )
```

Wait, `_DataBaseMetaClass` also raises an error if `fields` is `None` (which happens if `_FWD_DECL_FIELDS` is not in `attrs` and there is no `_proto_class`).

So if I remove the check in `_DataObjectBaseMetaClass`, `_DataBaseMetaClass` will still raise an error because `attrs.pop(mcs._FWD_DECL_FIELDS, None)` will return `None` for `Derived` if it doesn't have `__annotations__` (and thus `_DataObjectBaseMetaClass` didn't set `_FWD_DECL_FIELDS` in `attrs`).

So for `Derived`, we need to make sure that `_FWD_DECL_FIELDS` is present in `attrs` if it inherits from something that is a `DataObjectBase`.

Wait, if `Derived` inherits from `Base`, it should probably have the same fields as `Base`.

If `_DataObjectBaseMetaClass.__new__` sees no `__annotations__`, it should check if any base class has `_FWD_DECL_FIELDS`.

Wait, `_FWD_DECL_FIELDS` is set on the class itself. So `Base.__fwd_decl_fields__` should exist.

Let's check what `_DataBaseMetaClass._FWD_DECL_FIELDS` is. It's `"__fwd_decl_fields__"`.

So in `_DataObjectBaseMetaClass.__new__`:

```python
    def __new__(mcs, name, bases, attrs):
        # ...
        if name != "DataObjectBase":
            field_names = attrs.get("__annotations__")
            if field_names is None:
                # Check if any base class is a DataObjectBase and has fields
                # But wait, if it inherits from Base, Base already has the fields.
                # We just need to copy them to Derived so that _DataBaseMetaClass can find them.
                
                # Let's try to find a base class that has _FWD_DECL_FIELDS
                found_fields = None
                for base in bases:
                    # Use getattr because _FWD_DECL_FIELDS is an attribute name
                    fields = getattr(base, _DataBaseMetaClass._FWD_DECL_FIELDS, None)
                    if fields is not None:
                        found_fields = fields
                        break
                
                if found_fields is None:
                    raise TypeError(
                        "All DataObjectBase classes must follow dataclass syntax"
                    )
                field_names = found_fields

            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names

        # Delegate to the base metaclass
        return super().__new__(mcs, name, bases, attrs)
```

Wait, if `Derived` *does* have `__annotations__`, it should probably merge them with the base class's annotations? But the issue says "without adding or changing any of the schema". If it adds new fields, then it's adding to the schema.

Does `dataclasses.dataclass` merge annotations from base classes? Yes, it does.

However, `_DataBaseMetaClass` is trying to determine the fields *before* the class is fully created (in `__new__`).

If `Derived` adds new fields, `attrs.get("__annotations__")` will be present. Then `attrs[_DataBaseMetaClass._FWD_DECL_FIELDS]` will be set to those new fields.

But `_DataBaseMetaClass.__new__` then does:
```python
        fields = attrs.pop(mcs._FWD_DECL_FIELDS, None)
        # ...
        attrs["fields"] = fields
```
It only takes the fields from `_FWD_DECL_FIELDS`. If `Derived` has new fields and `Base` has old fields, and we only put `Derived`'s fields in `_FWD_DECL_FIELDS`, then `Derived` will only have those new fields in its `fields` attribute. That seems wrong.

Wait, let's check how `dataclass` works with inheritance.
If `Base` has `foo: int` and `Derived(Base)` has `bar: int`, `Derived` has both `foo` and `bar`.

If `caikit` is mimicking this, then `_DataObjectBaseMetaClass` should probably merge annotations.

But the bug report is specifically about the case where NO annotations are added.

Let's see if I can find where `_FWD_DECL_FIELDS` is used elsewhere.

I already did `grep -rn "_FWD_DECL_FIELDS" caikit/core/data_model/dataobject.py` and it was only there.
I should also check `caikit/core/data_model/base.py`. I already did, and it's used in `_DataBaseMetaClass.__new__` as I saw.

Wait, let's look at `_DataBaseMetaClass.__new__` again.

```python
        # Look for the set of fields either from a predefined protobuf class or
        # from a forward declaration from @dataobject
        fields = ()
        proto_class = None
        if name not in ["DataBase", "DataObjectBase"]:
            # ...
            # Otherwise, we need to get the fields from a "special" attribute
            else:
                fields = attrs.pop(mcs._FWD_DECL_FIELDS, None)
                # ...
```

If `Derived` inherits from `Base`, and `Base` had `_FWD_DECL_FIELDS` set, `Derived`'s `attrs` does NOT contain `_FWD_DECL_FIELDS` unless `_DataObjectBaseMetaClass` puts it there.

If `_DataObjectBaseMetaClass` puts `Base`'s fields into `Derived`'s `attrs[_FWD_DECL_FIELDS]`, then `_DataBaseMetaClass` will see them and set `Derived.fields = Base.fields`. This is correct if `Derived` didn't add new fields.

What if `Derived` DID add new fields?
`_DataObjectBaseMetaClass` sets `attrs[_FWD_DECL_FIELDS]` to `Derived`'s annotations.
Then `_DataBaseMetaClass` sets `Derived.fields` to `Derived`'s annotations.
But `Derived.fields` should be `Base.fields + Derived.fields`.

It seems `caikit` might have a bug with inheritance of fields in `DataObjectBase` classes anyway, or maybe it's intended that you can't add fields in subclasses?

Wait, the issue says: "overrides parent functions without adding or changing any of the schema". This implies the schema is inherited.

If I add fields to `Derived`, does it work?
Let's test.

{
  "command": "python3 - <<'PY'\nfrom caikit.core import DataObjectBase, dataobject\nfrom typing import Any\n\n@dataobject\nclass Base(DataObjectBase):\n    foo: int\n    def doit(self):\n        return self.foo * 2\n\n@dataobject\nclass Derived(Base):\n    bar: int\n    def doit(self):\n        return self.foo * 3\n\ntry:\n    b_inst = Base(1)\n    print(f\"Base instância: foo={b_inst.foo}\")\n    \n    d_inst = Derived(1, 2)\n    print(f\"Derived instância: foo={d_inst.foo}, bar={d_inst.bar}\")\n    print(\"Success\")\nexcept Exception as e:\n    print(f\"Error: {type(e).__name__}: {e}\")\nPY\n"
}

{
  "status": "ok",
  "stdout": "Base instância: foo=1\nDerived instância: foo=1, bar=2\nSuccess\n",
  "stderr": "<function register_backend_type at 0x7a75b9e51240> is still in the BETA phase and subject to change!\n",
  "exit_code": 0
}

Wait, it worked! How?

Let's look at `Derived` creation.
1. `_DataObjectBaseMetaClass.__new__` is called. `attrs` has `__annotations__ = {'bar': int}`.
2. It sets `attrs['_FWD_DECL_FIELDS'] = {'bar': int}`.
3. `_DataBaseMetaClass.__new__` is called. It pops `_FWD_DECL_FIELDS` and sets `Derived.fields = {'bar': int}`.
4. Then the class is created.
5. Then `@dataobject` decorator is called.
6. It calls `dataclasses.dataclass(Derived)`. `dataclass` merges annotations, so `Derived` has `foo` and `bar`.
7. Then it calls `_dataobject_to_proto(dataclass_=Derived)`.

Let's see `_dataobject_to_proto`.

{
  "end_line": 295,
  "filepath": "caikit/core/data_model/dataobject.py",
  "start_line": 269
}

{
  "status": "ok",
  "filepath": "caikit/core/data_model/dataobject.py",
  "content": "    def _is_python_optional(entry: Any) -> Any:\n        \"\"\"Detect if this type is a python optional\"\"\"\n        if get_origin(entry) is Union:\n            args = get_args(entry)\n            return type(None) in args\n\n\ndef _get_all_enums(\n    proto_class: Union[_message.Message, EnumTypeWrapper],\n) -> List[EnumTypeWrapper]:\n    \"\"\"Given a generated proto class, recursively extract all enums\"\"\"\n    all_enums = []\n    if isinstance(proto_class, EnumTypeWrapper):\n        all_enums.append(proto_class)\n    else:\n        for enum_descriptor in proto_class.DESCRIPTOR.enum_types:\n            all_enums.append(getattr(proto_class, enum_descriptor.name))\n        for nested_proto_descriptor in proto_class.DESCRIPTOR.nested_types:\n            all_enums.extend(\n                _get_all_enums(getattr(proto_class, nested_proto_descriptor.name))\n            )\n\n    return all_enums\n\n\ndef _make_data_model_class(proto_class, cls):\n    if issubclass(cls, DataObjectBase):",
  "start_line": 269,
  "end_line": 295,
  "total_lines": 386,
  "is_truncated": false
}

Wait, I haven't seen `_dataobject_to_proto`.

{
  "command": "grep -n \"def _dataobject_to_proto\" caikit/core/data_model/dataobject.py"
}

{
  "status": "ok",
  "stdout": "224:def _dataobject_to_proto(*args, **kwargs):\n",
  "stderr": "",
  "exit_code": 0
}

Wait, line 224. Let me read around there.

{
  "end_line": 240,
  "filepath": "caikit/core/data_model/dataobject.py",
  "start_line": 220
}

{
  "status": "ok",
  "filepath": "caikit/core/data_model/dataobject.py",
  "content": "\n## Implementation Details ######################################################\n\n\ndef _dataobject_to_proto(*args, **kwargs):\n    kwargs.setdefault(\"type_mapping\", DATAOBJECT_PY_TO_PROTO_TYPES)\n    return _DataobjectConverter(*args, **kwargs).descriptor\n\n\nclass _DataobjectConverter(DataclassConverter):\n    \"\"\"Augment the dataclass converter to be able to pull descriptors from\n    existing data objects\n    \"\"\"\n\n    def get_concrete_type(self, entry: Any) -> Any:\n        \"\"\"Also include data model classes and enums as concrete types\"\"\"\n        unwrapped = self._resolve_wrapped_type(entry)\n        if (\n            isinstance(unwrapped, type)\n            and issubclass(unwrapped, DataBase)\n            and unwrapped._proto_class is not None",
  "start_line": 220,
  "end_line": 240,
  "total_lines": 386,
  "is_truncated": false
}

`_dataobject_to_proto` uses `_DataobjectConverter`, which inherits from `DataclassConverter`. `DataclassConverter` is likely from some other module.

The key is that `@dataobject` calls `dataclasses.dataclass(cls)`, which handles the field inheritance.

And then `_dataobject_to_proto` uses the resulting dataclass to generate a proto.

So the only issue is the `TypeError` in `_DataObjectBaseMetaClass.__new__` when `__annotations__` is missing.

If I just change:

```python
            if field_names is None:
                raise TypeError(
                    "All DataObjectBase classes must follow dataclass syntax"
                )
```

to something that allows inheritance.

If `field_names` is `None`, we should check if any of the bases is a `DataObjectBase` (and not `DataObjectBase` itself, though `issubclass(base, DataObjectBase)` is what we want).

Wait, if we are in `_DataObjectBaseMetaClass.__new__`, and `name` is not `"DataObjectBase"`, then we are creating a subclass of `DataObjectBase`.

If `field_names` is `None`, we should check if any of the bases are `DataObjectBase` subclasses.

Wait, `DataObjectBase` itself is not a subclass of itself in the sense that it doesn't have `__annotations__`.

Let's check the bases.

```python
    def __new__(mcs, name, bases, attrs):
        # ...
        if name != "DataObjectBase":
            field_names = attrs.get("__annotations__")
            if field_names is None:
                # Check if any base is a subclass of DataObjectBase (and not DataObjectBase itself)
                # Actually, if ANY base is a subclass of DataObjectBase, it means it already has fields.
                # But we need to make sure we don't allow the first actual DataObject class to be empty.
                
                # a simple check:
                has_dataobject_base = any(
                    # We can't use issubclass(base, DataObjectBase) if DataObjectBase is not yet defined
                    # But we know that if a base is an instance of _DataObjectBaseMetaClass, it's a DataObjectBase-like class.
                    # Wait, DataObjectBase is also an instance of _DataObjectBaseMetaClass.
                    # But DataObjectBase doesn't have fields.
                    
                    # Maybe just check if any base has the _FWD_DECL_FIELDS attribute?
                     getattr(base, _DataBaseMetaClass._FWD_DECL_FIELDS, None) is not None
                    for base in bases
                )
                if not has_dataobject_base:
                    raise TypeError(
                        "All DataObjectBase classes must follow dataclass syntax"
                    )
                
                # If we found a base with fields, we should probably use those fields for now
                # so that _DataBaseMetaClass doesn't complain.
                for base in bases:
                    fields = getattr(base, _DataBaseMetaClass._FWD_DECL_FIELDS, None)
                    if fields is not None:
                        field_names = fields
                        break
            
            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names
```

Wait, if `field_names` is still `None` (which it will be if the loop didn't find anything), `attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = None` will be set, and then `_DataBaseMetaClass` will raise:
`error.value_check("<COR49310991E>", fields is not None, "No proto class found for {}", name)`

So if I use the logic:

```python
    def __new__(mcs, name, bases, attrs):
        # ...
        if name != "DataObjectBase":
            field_names = attrs.get("__annotations__")
            if field_names is None:
                # Look for fields in base classes
                for base in bases:
                    fields = getattr(base, _DataBaseMetaClass._FWD_DECL_FIELDS, None)
                    if fields is not None:
                        field_names = fields
                        break
                
                if field_names is None:
                    raise TypeError(
                        "All DataObjectBase classes must follow dataclass syntax"
                    )
            
            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names
```

This looks correct. If `Derived` has no annotations, it looks for them in its bases. If it finds some, it uses them. If it doesn't find any, it raises the `TypeError`.

One more thing: what if `Base` has annotations and `Derived` also has annotations?
Currently, `_DataObjectBaseMetaClass` does:
```python
            field_names = attrs.get("__annotations__")
            if field_names is None:
                # ... (my proposed fix)
            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names
```
If `Derived` has annotations, `field_names` will be `Derived`'s annotations. Then `attrs[_FWD_DECL_FIELDS]` is set to `Derived`'s annotations.
Then `_DataBaseMetaClass` sets `Derived.fields` to `Derived`'s annotations.

Is this correct? I thought `Derived.fields` should be `Base.fields + Derived.fields`.

Let's test this.

{
  "command": "python3 - <<'PY'\nfrom caikit.core import DataObjectBase, dataobject\nfrom typing import Any\n\n@dataobject\nclass Base(DataObjectBase):\n    foo: int\n    def doit(self):\n        return self.foo * 2\n\n@dataobject\nclass Derived(Base):\n    bar: int\n    def doit(self):\n        return self.foo * 3\n\nd_inst = Derived(1, 2)\nprint(f\"Derived fields: {d_inst.fields}\")\nPY\n"
}

{
  "status": "ok",
  "stdout": "Derived fields: ('foo', 'bar')\n",
  "stderr": "<function register_backend_type at 0x7f11a5d352d0> is still in the BETA phase and subject to change!\n",
  "exit_code": 0
}

Wait, `Derived.fields` is `('foo', 'bar')`! How?
I thought `_DataBaseMetaClass` did:
```python
        # Otherwise, we need to get the fields from a "special" attribute
        else:
            fields = attrs.pop(mcs._FWD_DECL_FIELDS, None)
            # ...
        attrs["fields"] = fields
```
And `_DataObjectBaseMetaClass` did:
```python
            field_names = attrs.get("__annotations__")
            # ...
            attrs[_DataBaseMetaClass._FWD_DECL_FIELDS] = field_names
```
So `attrs[_FWD_DECL_FIELDS]` should have been `{'bar': int}`. Thus `Derived.fields` should have been `('bar',)`?

Wait, let me check `_DataBaseMetaClass.__new__` in `caikit/core/data_model/base.py` more carefully.

```python
        # Look for the set of fields either from a predefined protobuf class or
        # from a forward declaration from @dataobject
        fields = ()
        proto_class = None
        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:
                # ... (extracts fields from proto_class)
                fields = ...
            # Otherwise, we need to get the fields from a "special" attribute
            else:
                fields = attrs.pop(mcs._FWD_DECL_FIELDS, None)
                # ...
                error.value_check(
                    "<COR49310991E>",
                    fields is not None,
                    "No proto class found for {}",
                    name,
                )
        attrs["fields"] = fields
```

Wait, if `fields` is a dictionary (from `__annotations__`), and then it's assigned to `attrs["fields"]`.

But in my test, `Derived.fields` was `('foo', 'bar')`.

Is it possible that `_DataBaseMetaClass.parse_proto_descriptor(instance)` is what sets `fields`?

Let's look at `_DataBaseMetaClass.parse_proto_descriptor`.

{
  "end_line": 990,
  "filepath": "caikit/core/data_model/base.py",
  "start_line": 900
}

{
  "status": "ok",
  "filepath": "caikit/core/data_model/base.py",
  "content": "        ):\n            return _recursive_to_dict(attr)\n\n        # fallback to the string representation\n        return str(attr)\n\n    @staticmethod\n    def get_class_for_proto(\n        proto: Union[Descriptor, ProtoMessageType]\n    ) -> Type[\"DataBase\"]:\n        \"\"\"Look up the data model class corresponding to the given protobuf\n\n        If no data model is found, this raises an AttributeError\n\n        Args:\n            proto (Union[Descriptor, ProtoMessageType])\n                The proto name or descriptor to look up against\n\n        Returns:\n            dm_class (Type[DataBase])\n                The data model class corresponding to the given protobuf\n        \"\"\"\n        error.type_check(\n            \"<COR46446770E>\",\n            Descriptor,\n            ProtoMessageType,\n            proto=proto,\n        )\n        proto_full_name = (\n            proto.full_name\n            if isinstance(proto, Descriptor)\n            else proto.DESCRIPTOR.full_name\n        )\n        cls = _DataBaseMetaClass.class_registry.get(proto_full_name)\n        if cls is None:\n            error(\n                \"<COR71783879E>\",\n                AttributeError(\n                    \"no data model class found in registry for protobufs named `{}`\".format(\n                        proto_full_name\n                    )\n                ),\n            )\n\n        return cls\n\n    @staticmethod\n    def get_class_for_name(class_name: str) -> Type[\"DataBase\"]:\n        \"\"\"Look up the data model class corresponding to the given name\n\n        This lookup attempts to encode various naming conventions that might be\n        used, but it can fail in multiple ways:\n\n        1. No class with the given name is known\n        2. Multiple classes with the same name, but different qualified parents\n           are found\n\n        A ValueError will be raised if either of the above happens\n\n        Args:\n            class_name (str)\n                The name of the class either as a fully-qualified protobuf name\n                or as the unqualified class name\n\n        Returns:\n            dm_class (Type[DataBase])\n                The data model class corresponding to the given protobuf\n        \"\"\"\n        dm_class = _DataBaseMetaClass.class_registry.get(class_name)\n        if dm_class is not None:\n            return dm_class\n        matching_classes = [\n            (full_name, dm_class)\n            for full_name, dm_class in _DataBaseMetaClass.class_registry.items()\n            if full_name.rpartition(\".\")[-1] == class_name\n        ]\n        if len(matching_classes) == 1:\n            return matching_classes[0][1]\n        if len(matching_classes) > 1:\n            error(\n                \"<COR02514290E>\",\n                ValueError(\n                    \"Conflicting data model classes for [{}]: {}\".format(\n                        class_name, [match[0] for match in matching_classes]\n                    )\n                ),\n            )\n        error(\n            \"<COR99562895E>\",\n            ValueError(f\"No data model class match for {class_name}\"),\n        )",
  "start_line": 900,
  "end_line": 990,
  "total_lines": 990,
  "is_truncated": false
}