feat: coerce nested dataclass fields in Relation.load

ops/charm.py

@@ -21,6 +21,8 @@
 import logging
 import os
 import pathlib
+import types
+import typing
 import warnings
 from collections.abc import Mapping
 from typing import (
@@ -33,6 +35,7 @@
     TypedDict,
     TypeVar,
     cast,
+    get_type_hints,
 )
 
 from . import model
@@ -1699,6 +1702,76 @@ def _juju_fields(cls: type[object]) -> dict[str, str]:
     raise ValueError('Unable to find class fields')
 
 
+def _coerce_field(tp: Any, value: Any) -> Any:
+    """Coerce a decoded ``value`` into the dataclass field type ``tp``.
+
+    Used by :meth:`ops.Relation.load` to recursively construct nested
+    dataclasses and enum values from JSON-decoded relation data. An
+    ``Optional``/``Union`` field is coerced against its single non-``None``
+    member; ``dict``/``Mapping`` fields are coerced against their value type;
+    a variable-length ``tuple[X, ...]`` is coerced element-wise against ``X``
+    and a fixed-length ``tuple[X, Y, ...]`` is coerced positionally.
+    """
+    origin = typing.get_origin(tp)
+    if origin is not None:
+        args = typing.get_args(tp)
+        if origin is typing.Union or origin is types.UnionType:
+            non_none = [a for a in args if a is not type(None)]
+            if len(non_none) == 1:
+                # Optional[X]: coerce against the one concrete member.
+                return _coerce_field(non_none[0], value)
+            # A Union of more than one concrete type: no way to tell which
+            # member to coerce against, so accept the value as-is.
+            return value
+        if origin is list and args:
+            return [_coerce_field(args[0], v) for v in value]
+        if origin is tuple and args:
+            if args[-1] is Ellipsis:
+                return tuple(_coerce_field(args[0], v) for v in value)
+            return tuple(_coerce_field(t, v) for t, v in zip(args, value, strict=True))
+        if origin in (set, frozenset) and args:
+            return {_coerce_field(args[0], v) for v in value}
+        if isinstance(origin, type) and issubclass(origin, Mapping) and len(args) == 2:
+            return {k: _coerce_field(args[1], v) for k, v in value.items()}
+        # Literal and other constructed generics: accept the value as-is.
+        return value
+    if isinstance(tp, type):
+        if dataclasses.is_dataclass(tp):
+            return _build_dataclass(tp, value)
+        if issubclass(tp, enum.Enum):
+            return tp(value)
+    return value
+
+
+def _build_dataclass(cls: Any, data: Mapping[str, Any], *args: Any) -> Any:
+    """Construct dataclass ``cls`` from ``data`` and any positional ``args``.
+
+    Recursively coerces nested dataclass / enum / list / set / tuple / dict
+    fields supplied via ``data``. Any leading fields already filled
+    positionally by ``args`` are matched by position, not by name, so they are
+    passed through as given rather than coerced.
+
+    Falls back to the un-coerced ``cls(*args, **data)`` if ``cls``'s type hints
+    can't be resolved, for example a ``TYPE_CHECKING``-only import with no
+    runtime name: ``get_type_hints`` resolves every field's annotation
+    eagerly, so one unresolvable field would otherwise break construction even
+    when the relation data at hand doesn't touch it.
+
+    Raises ``TypeError`` (via the dataclass ``__init__``) if a required field is
+    missing, and ``ValueError``/``TypeError`` from coercion of malformed values.
+    """
+    try:
+        hints = get_type_hints(cls)
+    except NameError:
+        return cls(*args, **data)
+    kwargs: dict[str, Any] = {}
+    for field in dataclasses.fields(cls)[len(args) :]:
+        if field.name not in data:
+            continue
+        kwargs[field.name] = _coerce_field(hints[field.name], data[field.name])
+    return cls(*args, **kwargs)
+
+
 class CharmMeta:
     """Object containing the metadata for the charm.
 

ops/model.py

@@ -1766,9 +1766,9 @@ def load(
     ) -> _T:
         """Load the data for this relation into an instance of a data class.
 
-        The raw Juju relation data is passed to the data class's ``__init__``
-        method as keyword arguments, with values decoded using the provided
-        decoder function, or :func:`json.loads` if no decoder is provided.
+        The raw Juju relation data is decoded using the provided decoder
+        function, or :func:`json.loads` if no decoder is provided, and passed
+        to the data class's ``__init__`` method as keyword arguments.
 
         For example::
 
@@ -1792,8 +1792,41 @@ def _observer(self, event: ops.RelationEvent):
                 data = event.relation.load(Data, event.app)
                 secret = self.model.get_secret(data.secret_id)
 
-        Any additional positional or keyword arguments will be passed through to
-        the data class ``__init__``.
+        For a Pydantic ``BaseModel`` or pydantic dataclass, the decoded values
+        are passed straight through as keyword arguments and Pydantic
+        performs its own coercion and validation.
+
+        For any other :func:`dataclasses.dataclass`, the decoded values are
+        also recursively coerced to match each field's type hint before being
+        passed to ``__init__``:
+
+        - A nested dataclass or :class:`enum.Enum` field is constructed from
+          its decoded value.
+        - ``list``, ``set``, and ``frozenset`` fields coerce each element
+          against the type argument.
+        - A variable-length ``tuple[X, ...]`` coerces every element against
+          ``X``; a fixed-length ``tuple[X, Y, ...]`` coerces each position
+          against its own type, and stays a ``tuple``.
+        - A ``dict``/``Mapping`` field coerces its values against the value
+          type.
+        - An ``Optional``/``Union`` field is coerced against its single
+          non-``None`` member; a ``Union`` of more than one concrete type is
+          passed through as-is, since there is no way to tell which member to
+          coerce against.
+        - ``Literal`` fields, and any other constructed generic not listed
+          above, are passed through unchanged.
+
+        If the class's type hints can't be resolved at all - for example, a
+        ``TYPE_CHECKING``-only import with no runtime name - the values are
+        passed through uncoerced instead of raising.
+
+        Any additional positional or keyword arguments will be passed through
+        to the data class ``__init__``. For a non-pydantic dataclass target,
+        positional arguments are matched to the class's leading fields by
+        position; those fields are passed through as given rather than
+        coerced, since there is no field name to coerce them against, but any
+        remaining fields supplied from the relation data are still coerced as
+        above.
 
         Args:
             cls: A class, typically a Pydantic `BaseModel` subclass or a
@@ -1823,6 +1856,15 @@ def _observer(self, event: ops.RelationEvent):
                 data[key] = decoder(value)
             elif key in fields:
                 data[fields[key]] = decoder(value)
+        # For plain (non-pydantic) dataclass targets, recursively coerce nested
+        # dataclass / enum / list / set fields. Pydantic handles its own coercion.
+        # '__pydantic_validator__' is what pydantic.dataclasses.is_pydantic_dataclass
+        # itself checks for; '__is_pydantic_dataclass__' only exists from pydantic
+        # 2.11, so relying on it misses every earlier 2.x pydantic dataclass.
+        # Any fields filled positionally by args are left uncoerced, since args
+        # are matched to the class's leading fields by position, not by name.
+        if dataclasses.is_dataclass(cls) and '__pydantic_validator__' not in cls.__dict__:
+            return _charm._build_dataclass(cls, data, *args)
         return cls(*args, **data)
 
     def save(

test/test_model_relation_data_class.py

@@ -21,10 +21,16 @@
 import json
 import urllib.parse
 from collections.abc import Callable, Iterable
-from typing import Any, Protocol, cast
+from typing import TYPE_CHECKING, Any, Protocol, cast
 
 import pytest
 
+if TYPE_CHECKING:
+    # Used only by test_relation_load_falls_back_when_type_hints_unresolvable,
+    # which needs an annotation naming a type that is never actually imported
+    # at runtime.
+    import decimal
+
 try:
     import pydantic
     import pydantic.dataclasses
@@ -45,6 +51,11 @@ class Nested:
     sub: int = 28
 
 
+class _Colour(enum.Enum):
+    RED = 'red'
+    BLUE = 'blue'
+
+
 class DatabagProtocol(Protocol):
     foo: str
     bar: int
@@ -432,6 +443,168 @@ def _on_relation_changed(self, event: ops.RelationChangedEvent):
     assert obj.c == 'foo'
 
 
+def _load_into(cls: type[Any], remote_app_data: dict[str, str]) -> Any:
+    """Load ``remote_app_data`` into ``cls`` via ``Relation.load`` and return the result."""
+
+    class Charm(ops.CharmBase):
+        def __init__(self, framework: ops.Framework):
+            super().__init__(framework)
+            framework.observe(self.on['db'].relation_changed, self._on_relation_changed)
+
+        def _on_relation_changed(self, event: ops.RelationChangedEvent):
+            self.data = event.relation.load(cls, event.app)
+
+    ctx = testing.Context(Charm, meta={'name': 'foo', 'requires': {'db': {'interface': 'db-int'}}})
+    rel = testing.Relation('db', remote_app_data=remote_app_data)
+    state_in = testing.State(leader=True, relations={rel})
+    with ctx(ctx.on.relation_changed(rel), state_in) as mgr:
+        mgr.run()
+        return mgr.charm.data
+
+
+def test_relation_load_optional_nested_dataclass():
+    """Optional[X] (a Union with one concrete member) is coerced against X."""
+
+    @dataclasses.dataclass
+    class Data:
+        inner: Nested | None = None
+
+    obj = _load_into(Data, {'inner': json.dumps({'sub': 1})})
+    assert isinstance(obj.inner, Nested)
+    assert obj.inner.sub == 1
+
+    obj = _load_into(Data, {})
+    assert obj.inner is None
+
+
+def test_relation_load_dict_of_nested_dataclass():
+    """dict[str, X] fields are coerced against X for each value."""
+
+    @dataclasses.dataclass
+    class Data:
+        by_name: dict[str, Nested]
+
+    obj = _load_into(Data, {'by_name': json.dumps({'a': {'sub': 1}, 'b': {'sub': 2}})})
+    assert obj.by_name == {'a': Nested(sub=1), 'b': Nested(sub=2)}
+    assert all(isinstance(v, Nested) for v in obj.by_name.values())
+
+
+def test_relation_load_union_of_two_concrete_types_passes_through():
+    """A Union with more than one concrete member is passed through as-is.
+
+    There is no way to tell which member to coerce against, so this is a
+    regression check that such fields keep working uncoerced rather than
+    raising.
+    """
+
+    @dataclasses.dataclass
+    class Data:
+        value: int | str
+
+    obj = _load_into(Data, {'value': json.dumps('x')})
+    assert obj.value == 'x'
+
+
+def test_relation_load_variable_length_tuple():
+    """tuple[X, ...] is coerced element-wise against X and stays a tuple."""
+
+    @dataclasses.dataclass
+    class Data:
+        items: tuple[Nested, ...]
+
+    obj = _load_into(Data, {'items': json.dumps([{'sub': 1}, {'sub': 2}])})
+    assert obj.items == (Nested(sub=1), Nested(sub=2))
+    assert isinstance(obj.items, tuple)
+
+
+def test_relation_load_heterogeneous_tuple():
+    """A fixed-length tuple[X, Y] is coerced positionally against each type."""
+
+    @dataclasses.dataclass
+    class Data:
+        pair: tuple[int, _Colour]
+
+    obj = _load_into(Data, {'pair': json.dumps([1, 'red'])})
+    assert obj.pair == (1, _Colour.RED)
+    assert isinstance(obj.pair, tuple)
+
+
+def test_relation_load_pydantic_dataclass_guard_without_is_pydantic_dataclass():
+    """The pydantic guard must key off __pydantic_validator__, not __is_pydantic_dataclass__.
+
+    __is_pydantic_dataclass__ only exists from pydantic 2.11; older pydantic
+    dataclasses (as old as 2.0.3) have __pydantic_validator__ in their
+    __dict__ instead. Simulate that older shape on a plain dataclass, without
+    needing multiple installed pydantic versions, and confirm Relation.load
+    still treats it as a pydantic target: ops's own recursive coercion must
+    not run, so a nested-dataclass-typed field stays a plain decoded dict
+    rather than being (mis-)coerced ahead of pydantic's own validation.
+    """
+
+    @dataclasses.dataclass
+    class Data:
+        nested: Nested
+
+    # Simulate pydantic < 2.11's shape.
+    Data.__pydantic_validator__ = object()  # pyright: ignore[reportAttributeAccessIssue]
+
+    obj = _load_into(Data, {'nested': json.dumps({'sub': 1})})
+    assert isinstance(obj.nested, dict)
+
+
+def test_relation_load_falls_back_when_type_hints_unresolvable():
+    """get_type_hints raises NameError on a TYPE_CHECKING-only annotation.
+
+    ops's own ruff config disables TC001/2/3, so charms following ops's
+    conventions are the ones most likely to hit this. Relation.load must
+    fall back to the un-coerced constructor rather than raising, matching
+    what main's cls(**data) path already did before recursive coercion
+    existed.
+    """
+
+    @dataclasses.dataclass
+    class Data:
+        amount: decimal.Decimal | None = None
+        name: str = ''
+
+    obj = _load_into(Data, {'name': json.dumps('x')})
+    assert obj.name == 'x'
+    assert obj.amount is None
+
+
+def test_relation_load_extra_args_still_coerces_remaining_fields():
+    """A positional arg must not silently disable coercion for other fields.
+
+    relation.load(cls, src, *args) matches args to cls's leading fields by
+    position; any fields filled that way are left uncoerced (there's nothing
+    to coerce them against without knowing which field each arg is for), but
+    fields still supplied from the relation data should keep being coerced.
+    """
+
+    @dataclasses.dataclass
+    class Data:
+        a: int
+        b: Nested
+
+    class Charm(ops.CharmBase):
+        def __init__(self, framework: ops.Framework):
+            super().__init__(framework)
+            framework.observe(self.on['db'].relation_changed, self._on_relation_changed)
+
+        def _on_relation_changed(self, event: ops.RelationChangedEvent):
+            self.data = event.relation.load(Data, event.app, 10)
+
+    ctx = testing.Context(Charm, meta={'name': 'foo', 'requires': {'db': {'interface': 'db-int'}}})
+    rel = testing.Relation('db', remote_app_data={'b': json.dumps({'sub': 1})})
+    state_in = testing.State(leader=True, relations={rel})
+    with ctx(ctx.on.relation_changed(rel), state_in) as mgr:
+        mgr.run()
+        obj = mgr.charm.data
+    assert obj.a == 10
+    assert isinstance(obj.b, Nested)
+    assert obj.b.sub == 1
+
+
 @pytest.mark.parametrize('charm_class', _test_classes)
 def test_relation_save_simple(charm_class: type[BaseTestCharm]):
     class Charm(charm_class):