Polymorphic (de)serializastion support (#7104)

Author: hannes-ucsc

src/azul/attrs.py

@@ -41,6 +41,7 @@
     require,
 )
 from azul.json import (
+    PolymorphicSerializable,
     Serializable,
 )
 from azul.types import (
@@ -483,6 +484,10 @@ def _metadata[V](self, key: str, default: V) -> V:
             except KeyError:
                 return default
 
+        @cached_property
+        def discriminator(self) -> str | None:
+            return self._metadata('discriminator', None)
+
         def handle(self, x: str) -> T:
             if self.custom is None:
                 return self._handle(x, self._reify(self.field.type))
@@ -536,7 +541,12 @@ def _handle(self, x: str, field_type: Any):
                 elif issubclass(field_type, Serializable):
                     inner_cls_name = field_type.__name__
                     self.globals[inner_cls_name] = field_type
-                    return self._serializable(x, inner_cls_name)
+                    is_polymorphic = issubclass(field_type, PolymorphicSerializable)
+                    has_discriminator = self.discriminator is not None
+                    if is_polymorphic and has_discriminator:
+                        return self._polymorphic(x, inner_cls_name)
+                    else:
+                        return self._serializable(x, inner_cls_name)
             else:
                 origin = get_origin(field_type)
                 if origin in (Union, UnionType):
@@ -575,6 +585,10 @@ def _optional(self, x: str, field_type: type) -> T:
         def _serializable(self, x: str, inner_cls_name: str) -> T:
             raise NotImplementedError
 
+        @abstractmethod
+        def _polymorphic(self, x: str, inner_cls_name: str) -> T:
+            raise NotImplementedError
+
         @abstractmethod
         def _list(self, x: str, item_type: type) -> T:
             raise NotImplementedError
@@ -603,6 +617,15 @@ def _serializable(self, x: str, inner_cls_name: str) -> Source:
                 f'{x} = {inner_cls_name}.from_json({x})'
             ]
 
+        def _polymorphic(self, x: str, inner_cls_name: str) -> Source:
+            depth = next(self.depth)
+            cls = f'cls{depth}'
+            return [
+                f'{cls} = {x}["{self.discriminator}"]',
+                f'{cls} = {inner_cls_name}.cls_from_json({cls})',
+                f'{x} = {cls}.from_json({x})'
+            ]
+
         def _primitive(self, x: str, field_type: type) -> Source:
             return [
                 f'if not isinstance({x}, {field_type.__name__}):', [
@@ -670,6 +693,9 @@ def _optional(self, x: str, field_type: type) -> str:
         def _serializable(self, x: str, inner_cls_name: str) -> str:
             return f'{x}.to_json()'
 
+        def _polymorphic(self, x: str, inner_cls_name: str) -> str:
+            return f'dict({x}.to_json(), {self.discriminator}={x}.cls_to_json())'
+
         def _list(self, x: str, item_type: type) -> str:
             depth = next(self.depth)
             v = f'v{depth}'
@@ -740,3 +766,95 @@ def _set_field_metadata[T: attrs.Attribute](field: T | None, key, value):
     metadata = field.metadata.setdefault('azul', {})
     metadata[key] = value
     return field
+
+
+def polymorphic[T: attrs.Attribute](field: T | None = None,
+                                    *,
+                                    discriminator: str
+                                    ) -> T:
+    """
+    Mark an attrs field to use the given name for the discriminator property in
+    serialized instances of PolymorphicSerializable that occur in the value of
+    that field. The given discriminator property of a serialized instance
+    represents the type to use when deserializing that instance again.
+
+    >>> from azul.json import RegisteredPolymorphicSerializable
+
+    >>> class Inner(SerializableAttrs, RegisteredPolymorphicSerializable):
+    ...     pass
+
+    >>> @attrs.frozen
+    ... class InnerWithInt(Inner):
+    ...     x: int
+
+    >>> @attrs.frozen
+    ... class InnerWithStr(Inner):
+    ...     y: str
+
+    >>> @attrs.frozen(kw_only=True)
+    ... class Outer(SerializableAttrs):
+    ...     inner: Inner = polymorphic(discriminator='type')
+    ...     inners: list[Inner] = polymorphic(discriminator='_cls')
+
+    >>> from azul.doctests import assert_json
+
+    >>> outer = Outer(inner=InnerWithInt(42),
+    ...               inners=[InnerWithStr('foo'), InnerWithInt(7)])
+    >>> assert_json(outer.to_json())
+    {
+        "inner": {
+            "x": 42,
+            "type": "InnerWithInt"
+        },
+        "inners": [
+            {
+                "y": "foo",
+                "_cls": "InnerWithStr"
+            },
+            {
+                "x": 7,
+                "_cls": "InnerWithInt"
+            }
+        ]
+    }
+    >>> Outer.from_json(outer.to_json()) == outer
+    True
+
+    In order to enable polymorphic serialization of the value of a given field,
+    the discriminator property needs to be specified explicitly, otherwise the
+    serialization framework will resort to the static type of the field.
+
+    >>> @attrs.frozen
+    ... class GenericOuter[T: Inner](SerializableAttrs):
+    ...     inner: T
+
+    >>> class StaticOuter(GenericOuter[InnerWithInt]):
+    ...     pass
+
+    >>> outer = StaticOuter(InnerWithInt(42))
+    >>> outer.to_json()
+    {'inner': {'x': 42}}
+
+    Despite the fact that ``{'x': 42}`` does not encode any type information,
+    ``from_json`` can tell from the static type of the field that {'x': 42}
+    should be deserialized as an ``InnerWithInt``.
+
+    >>> StaticOuter.from_json(outer.to_json()).inner
+    InnerWithInt(x=42)
+
+    >>> StaticOuter.from_json(outer.to_json()) == outer
+    True
+
+    However, when the static type of the field is not concrete, deserialization
+    may fail or, like in this case, lose information by creating an instance of
+    the parent class instead of the class that was serialized.
+
+    >>> @attrs.frozen
+    ... class AbstractOuter(SerializableAttrs):
+    ...     inner: Inner
+
+    >>> outer = AbstractOuter(InnerWithInt(42))
+    >>> AbstractOuter.from_json(outer.to_json()).inner  # doctest: +ELLIPSIS
+    <azul.attrs.Inner object at ...>
+    """
+    return _set_field_metadata(field, 'discriminator', discriminator)

src/azul/json.py

@@ -267,6 +267,67 @@ def to_json(self) -> AnyJSON:
         raise NotImplementedError
 
 
+class PolymorphicSerializable(Serializable):
+    """
+    A class whose subclasses' instances can be transformed to and from JSON
+    while retaining the concrete type of said instances.
+    """
+
+    @classmethod
+    def cls_to_json(cls) -> AnyJSON:
+        """
+        Serialize the given type to JSON.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def cls_from_json(cls, json: AnyJSON) -> type[Self]:
+        """
+        Deserialize a subtype of the given type from the given JSON.
+        """
+        raise NotImplementedError
+
+
+class RegisteredPolymorphicSerializable(PolymorphicSerializable):
+    """
+    A polymorphically serializable class that tracks its subclasses in a
+    registry and uses their name to discriminate serialized instances. It
+    requires every subclass to be registered before instances of that subclass
+    can be (de)serialized. It also requires the name of each subclass to be
+    unique, regardless of the module the subclass is defined in.
+    """
+
+    _registry: dict[str, type[Self]] = {}
+
+    @classmethod
+    def cls_to_json(cls) -> AnyJSON:
+        assert cls._registry[cls.__name__] == cls
+        return cls.__name__
+
+    @classmethod
+    def cls_from_json(cls, json: AnyJSON) -> type[Self]:
+        return cls._registry[json_str(json)]
+
+    def __init_subclass__(cls):
+        super().__init_subclass__()
+        try:
+            other_cls = cls._registry[cls.__name__]
+        except KeyError:
+            pass
+        else:
+            # For attrs classes, this hook is invoked twice: once for the
+            # original class and once for the attrs-generated replacement. These
+            # are two different objects, so they are neither the same nor equal
+            # so it is difficult to tell wether we're dealing with the attrs
+            # replacement or a genuine collision. Both original and replacement
+            # reference the same containing module, so we assume that two
+            # classes of the same name from the same module indicate that attrs
+            # is involved and does not constitue a collision.
+            assert other_cls.__module__ == cls.__module__, R(
+                'Class name collision', cls, other_cls)
+        cls._registry[cls.__name__] = cls
+
+
 class Parseable(Serializable):
     """
     A class whose instances have a string representation that can be used in