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
@@ -1680,10 +1683,13 @@ def _juju_fields(cls: type[object]) -> dict[str, str]:
         for field in dataclasses.fields(cls):
             alias = field.metadata.get('alias', field.name)
             # If this a Pydantic dataclass, then it handles the alias.
-            # Using pydantic.dataclasses.is_pydantic_dataclass() would be
-            # best here, but we don't want to import pydantic in ops, so
-            # we look more explicitly.
-            if getattr(cls, '__is_pydantic_dataclass__', False):
+            # Using pydantic.dataclasses.is_pydantic_dataclass() would be best
+            # here, but we don't want to import pydantic in ops, so we check
+            # for the attribute that function itself checks for. Note that
+            # '__is_pydantic_dataclass__' only exists from pydantic 2.11, so
+            # relying on that one misses every earlier 2.x pydantic dataclass,
+            # which reads the aliases back under their field names instead.
+            if '__pydantic_validator__' in cls.__dict__:
                 juju_to_arg[alias] = alias
             else:
                 juju_to_arg[alias] = field.name
@@ -1699,6 +1705,115 @@ 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.
+
+    Raises ``TypeError`` if the value for a sequence field is a string, bytes
+    or a mapping, or if the value for a mapping field is not a mapping: those
+    are all iterable, so coercing them element-wise would quietly produce a
+    wrong answer rather than fail.
+    """
+    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, unless
+                # the value really is None, which X itself won't accept.
+                if value is None:
+                    return None
+                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
+        # A str, bytes or mapping is iterable, so coercing element-wise would
+        # silently succeed with nonsense: a list of characters, or of the
+        # mapping's keys. None of those is a sequence the charm meant, so
+        # refuse rather than hand back the wrong answer.
+        if (
+            origin in (list, tuple, set, frozenset)
+            and args
+            and isinstance(value, (str, bytes, Mapping))
+        ):
+            given = cast('Any', value)
+            raise TypeError(f'expected a sequence for {tp}, got {type(given).__name__}: {given!r}')
+        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 is set and args:
+            return {_coerce_field(args[0], v) for v in value}
+        if origin is frozenset and args:
+            return frozenset(_coerce_field(args[0], v) for v in value)
+        if isinstance(origin, type) and issubclass(origin, Mapping) and len(args) == 2:
+            if not isinstance(value, Mapping):
+                raise TypeError(
+                    f'expected a mapping for {tp}, got {type(value).__name__}: {value!r}'
+                )
+            mapping = cast('Mapping[Any, Any]', value)
+            return {k: _coerce_field(args[1], v) for k, v in mapping.items()}
+        # Literal and other constructed generics: accept the value as-is.
+        return value
+    if isinstance(tp, type):
+        if dataclasses.is_dataclass(tp):
+            if isinstance(value, tp):
+                # Already the class we want: a caller's keyword argument passed
+                # through to `Relation.load`, rather than anything from a
+                # databag, so there is nothing to coerce.
+                return value
+            if not isinstance(value, Mapping):
+                # Without this, a remote app writing a string or a list where a
+                # nested dataclass belongs gets a default-constructed object
+                # that corresponds to nothing in the databag: `field.name not in
+                # 'oops'` is False for every field, so every one is skipped.
+                raise TypeError(
+                    f'expected a mapping for {tp.__name__}, got {type(value).__name__}: {value!r}'
+                )
+            return _build_dataclass(tp, cast('Mapping[str, Any]', 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(