How to expose tagged union and alternative constructors in GraphQL

[thomascobb]

I have a module called scanspec which allows the definition of a trajectory of points that will be moved through to collect experimental data. It is currently written with pydantic, but the Union of subclasses pattern introduces some messiness, and I would prefer to use plan dataclasses:
https://github.com/dls-controls/scanspec

Here is the core of what I would like to achieve (this doesn't currently work with apischema)

from typing import AsyncIterable, TypeVar

from dataclasses import dataclass

from typing_extensions import Annotated as A
from apischema import schema, deserialize, serialize


def desc(description: str):
    return schema(description=description)


class ScanSpec:
    async def points(self) -> AsyncIterable[float]:
        """Iterate through the points of the scan"""
        raise NotImplementedError(self)


T = TypeVar("T")


@dataclass
class Line(ScanSpec):
    """A straight line"""
    # Only declare descriptions once
    _start = desc("The first point")
    start: A[float, _start]
    _stop = desc("The last point")
    stop: A[float, _stop]
    step: A[float, desc("The step between points")] = 1

    async def points(self) -> AsyncIterable[float]:
        for point in range(self.start, self.stop, self.step):
            yield point

    @classmethod
    def sized(
        cls: T,
        start: A[float, _start],
        stop: A[float, _stop],
        size: A[int, desc("Number of points")],
    ) -> T:
        """Alternative constructor with size instead of step"""
        return cls(start, stop, stop - start / size)


@dataclass
class Concat(ScanSpec):
    left: A[ScanSpec, desc("First spec to produce")]
    right: A[ScanSpec, desc("Second spec to produce")]

    async def points(self) -> AsyncIterable[float]:
        async for point in self.left.points():
            yield point
        async for point in self.right.points():
            yield point


async def get_points(spec: ScanSpec) -> AsyncIterable[float]:
    async for point in spec.points():
        yield point


def test_de_serialize():
    line = Line(1, 4)
    serialized = {"Line": {"start": 1, "stop": 4, "step": 1}}
    assert serialize(line) == serialized
    assert deserialize(ScanSpec, serialized) == line

I would like to expose a GraphQL API that exposes a subscription get_points and queries for the constructors and alternative constructors of all ScanSpec subclasses that returns the serialized data shown. I was going to write this myself, but it looks like apischema may provide me some useful tools to do this.

As per discussion on graphql-python/graphene#729 (comment)

[wyfo]

Ok, what you want is in fact quite more complex than the [class as union of its subclasses example] of apischema documentation.
While in the example, you add deserializers to the base class but let the subclasses untouched, in your case, you have to modify the serialization of subclasses. And it's a complex serialization, because it's kind of recursive, but you have to prevent this recursion to happen.

More detailed, if you consider you have a created a tagged union type (let's name it ScanSpecTaggedUnion, I will detail this part after), you want Line serializing to ScanSpecTaggedUnion, but the Line field of ScanSpecTaggedUnion be serialized as a normal dataclass. That's 2 different serializations for the same type.

This is not a real issue, as apischema provides the material to do it, but the solution will require a little bit more code, and actually, there is two solutions:

• use a dynamic conversions ScanSpec --> ScanSpecTaggedUnion, inherited by the subclasses, while the Line field is serialized as a normal dataclass; it requires however to apply this conversion on each resolver/operation where you need it (using conversions parameter of resolver/Operation);
• register a serializer ScanSpec --> ScanSpecTaggedUnion (which will be inherited by its subclasses), but also add an identity field serializer with a sub-conversion bypassing the registered serialization; by the way, if you want to have a resolver serializing a Line without getting a ScanSpecTaggedUnion, you will have to use a dynamic conversion for this resolver too.

Personally, I prefer the first solution as I consider that Line could be also used without conversion in an other resolver, so I will use it for the rest of my answer, but I can also give you an example with the second.

There is again two possibilities of implementation:

• using __init_subclasses__ to automatically update all the conversions stuff when a new type is added;
• create conversions stuff all at once using __subclasses__.

I will arbitrary choose the second one, so it will give something like:

from dataclasses import dataclass, fields, make_dataclass
from typing import Any, Union

from graphql.utilities import print_schema

from apischema import (
    Undefined,
    UndefinedType,
    deserialize,
    deserializer,
    serialize,
    validator,
)
from apischema.conversions import Conversion
from apischema.graphql import Query, graphql_schema
from apischema.json_schema import deserialization_schema, serialization_schema


# Waiting for the tagged union feature, let's define something that come close to it
@dataclass
class _TaggedUnion:
    @validator
    def _get_only_field(self) -> Any:
        defined_fields = [
            f.name for f in fields(self) if getattr(self, f.name) is not Undefined
        ]
        if len(defined_fields) != 1:
            raise ValueError("Tagged union must have one and only one field set")
        return getattr(self, defined_fields[0])


class Base:
    pass


@dataclass
class Foo(Base):
    foo: int


@dataclass
class Bar(Base):
    bar: str


# Define the BaseTaggedUnion class
tagged_fields = [
    (sub_cls.__name__, Union[sub_cls, UndefinedType], Undefined)
    for sub_cls in Base.__subclasses__()  # no need to recurse in your example
]
tagged_union = make_dataclass(f"BaseTaggedUnion", tagged_fields, bases=(_TaggedUnion,))
# Conversions
deserializer(Conversion(_TaggedUnion._get_only_field, source=tagged_union, target=Base))
tagged_union_conversion = Conversion(
    lambda obj: tagged_union(**{type(obj).__name__: obj}),
    source=Base,
    target=tagged_union,
)

assert serialize(Foo(0), conversions=tagged_union_conversion) == {"Foo": {"foo": 0}}
assert deserialize(Base, {"Bar": {"bar": "hello"}}) == Bar("hello")
assert deserialization_schema(Base) == {
    "type": "object",
    "properties": {
        "Foo": {
            "type": "object",
            "properties": {"foo": {"type": "integer"}},
            "required": ["foo"],
            "additionalProperties": False,
        },
        "Bar": {
            "type": "object",
            "properties": {"bar": {"type": "string"}},
            "required": ["bar"],
            "additionalProperties": False,
        },
    },
    "additionalProperties": False,
    "$schema": "http://json-schema.org/draft/2019-09/schema#",
}
assert serialization_schema(Base, conversions=tagged_union_conversion) == {
    "type": "object",
    "properties": {
        "Foo": {
            "type": "object",
            "properties": {"foo": {"type": "integer"}},
            "required": ["foo"],
            "additionalProperties": False,
        },
        "Bar": {
            "type": "object",
            "properties": {"bar": {"type": "string"}},
            "required": ["bar"],
            "additionalProperties": False,
        },
    },
    "additionalProperties": False,
    "$schema": "http://json-schema.org/draft/2019-09/schema#",
}


def query(base: Base) -> Base:
    return base


schema = graphql_schema(query=[Query(query, conversions=tagged_union_conversion)])
schema_str = """\
type Query {
  query(base: BaseTaggedUnionInput!): BaseTaggedUnion!
}

type BaseTaggedUnion {
  Foo: Foo
  Bar: Bar
}

type Foo {
  foo: Int!
}

type Bar {
  bar: String!
}

input BaseTaggedUnionInput {
  Foo: FooInput
  Bar: BarInput
}

input FooInput {
  foo: Int!
}

input BarInput {
  bar: String!
}
"""
assert print_schema(schema) == schema_str
query_str = """
{
    query(base: {Foo: {foo: 0}}){
        Foo {
            foo
        }
        Bar {
            bar
        }
    }
}
"""
assert graphql.graphql_sync(schema, query_str).data == {
    "query": {
        "Foo": {"foo": 0},
        "Bar": None,
    }
}

Of course, BaseTaggedUnionInput is not a real tagged union because GraphQL still doesn't define them, but its validator (_get_only_field) ensure that only one field can be passed.

By the way, I invite you to update your apischema version to the last one (0.14.1), because I have done some refactoring thanks to your use case (especially concerning subscription where conversions were not supported yet).

[thomascobb]

Thanks for the detailed response, I'll try and get my head round this and get back to you...

[wyfo]

EDIT
An example of this use case is now presented in the documentation.

Actually, I've realized I've missed the whole point on alternative constructors … By the way, I've built my previous example on a symmetry between ScanSpec input and output, but you seemed to be only interested in the input part, so it's another misunderstanding on my side.

I will try to be more in line with your issue with this second response. For that, I will use the new release of apischema which brings the following elements : dataclass_input_wrapper and TaggedUnion. With both of them, things become a lot easier.

Here is your original example reworked with apischema:

import asyncio
from collections.abc import Iterator
from dataclasses import dataclass, field
from types import new_class
from typing import (
    Annotated,
    AsyncIterable,
    Callable,
    ClassVar,
    Collection,
    TypeVar,
)

import graphql

from apischema.conversions import (
    Conversion,
    dataclass_input_wrapper,
    reset_deserializers,
)
from apischema.graphql import graphql_schema
from apischema.json_schema import deserialization_schema
from apischema.tagged_unions import Tagged, TaggedUnion, get_tagged

T = TypeVar("T")

# Recursive implementation of type.__subclasses__
def rec_subclasses(cls: type[T]) -> Iterator[type[T]]:
    for sub_cls in cls.__subclasses__():
        yield sub_cls
        yield from rec_subclasses(sub_cls)


# Shortcut
def desc(description: str):
    return schema(description=description)


class ScanSpec:
    async def points(self) -> AsyncIterable[float]:
        """Iterate through the points of the scan"""
        raise NotImplementedError

    # Will be used in in the input tagged union
    _additional_constructors: ClassVar[Collection[Callable[..., "ScanSpec"]]] = []

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        namespace, annotations = {}, {}
        for spec_cls in rec_subclasses(ScanSpec):
            # Add tagged field for the ScanSpec subclass
            annotations[spec_cls.__name__] = Tagged[spec_cls]
            # Add tagged fields for all its additional constructors
            # (use class __dict__ in order to avoid inheritances of this constructors)
            for constructor in spec_cls.__dict__.get("_additional_constructors", ()):
                # Deref the constructor function if it is a classmethod/staticmethod
                constructor = constructor.__get__(None, spec_cls)
                # Build the alias of the field
                alias = "".join(map(str.capitalize, constructor.__name__.split("_")))
                # dataclass_input_wrapper uses get_type_hints, but the constructor
                # return type is stringified and the class not defined yet,
                # so it must be assigned manually
                constructor.__annotations__["return"] = spec_cls
                # Wraps the constructor and rename its input class
                wrapper, wrapper_cls = dataclass_input_wrapper(constructor)
                wrapper_cls.__name__ = alias
                # Add constructor tagged field with its conversion
                annotations[alias] = Tagged[spec_cls]
                namespace[alias] = Tagged(deserialization=wrapper)
        # Create the tagged union class
        namespace |= {"__annotations__": annotations}
        tagged_union = new_class(
            "ScanSpec",
            (TaggedUnion,),
            exec_body=lambda ns: ns.update(namespace),
        )
        # Because deserializers stack, they must be reset before being reassigned
        reset_deserializers(ScanSpec)
        # Register the deserializer using get_tagged
        deserializer(Conversion(lambda obj: get_tagged(obj)[1], tagged_union, ScanSpec))


@dataclass
class Line(ScanSpec):
    """A straight line"""

    start: float = field(metadata=desc("The first point"))
    stop: float = field(metadata=desc("The last point"))
    step: float = field(
        default=1, metadata=schema(description="The step between points", exc_min=0)
    )

    async def points(self) -> AsyncIterable[float]:
        point = self.start
        while point <= self.stop:
            yield point
            point += self.step

    @staticmethod
    def sized_line(
        start: Annotated[float, desc("The first point")],
        stop: Annotated[float, desc("The last point")],
        size: Annotated[float, schema(description="Number of points", min=1)],
    ) -> "Line":
        """Alternative constructor with size instead of step"""
        return Line(start=start, stop=stop, step=(stop - start) / (size - 1))

    _additional_constructors = [sized_line]


@dataclass
class Concat(ScanSpec):
    left: ScanSpec = field(metadata=desc("First spec to produce"))
    right: ScanSpec = field(metadata=desc("Second spec to produce"))

    async def points(self) -> AsyncIterable[float]:
        async for point in self.left.points():
            yield point
        async for point in self.right.points():
            yield point


async def get_points(spec: ScanSpec) -> AsyncIterable[float]:
    async for point in spec.points():
        yield point


def hello() -> str:
    return "world!"


scan_spec_schema = graphql_schema(query=[hello], subscription=[get_points])
assert (
    graphql.utilities.print_schema(scan_spec_schema)
    == '''\
type Query {
  hello: String!
}

type Subscription {
  getPoints(spec: ScanSpecInput!): Float!
}

input ScanSpecInput {
  Line: LineInput
  SizedLine: SizedLineInput
  Concat: ConcatInput
}

input LineInput {
  """The first point"""
  start: Float!

  """The last point"""
  stop: Float!

  """The step between points"""
  step: Float! = 1
}

input SizedLineInput {
  """The first point"""
  start: Float!

  """The last point"""
  stop: Float!

  """Number of points"""
  size: Float!
}

input ConcatInput {
  """First spec to produce"""
  left: ScanSpecInput!

  """Second spec to produce"""
  right: ScanSpecInput!
}
'''
)

subscription_query = """
subscription {
    getPoints(spec: {SizedLine: {start: 0, stop: 12, size: 3}})
}
"""


async def main():
    subscription = await graphql.subscribe(
        scan_spec_schema, graphql.parse(subscription_query)
    )
    assert [res.data["getPoints"] async for res in subscription] == [0.0, 6.0, 12.0]


asyncio.run(main())

This time, I've used __init_subclasses__, but its code could be simply executed once before the schema generation and it would works the same.

[wyfo]

Yes, dynamic conversions apply only to the first level (in order to ensure consistency and reuse of subschemas as explained here). Propagate them to class fields is a little tedious (but amelioration is coming), but here, i think default conversion is more suited for your need, both for deserialization and serialization.

(Also, thanks to your use case, I've found a little bug in the handling of schema type name when conversions are invovled, so I invite you to update your apischema version to run the following examples. You may notice than bypassing registered conversion is now done using directly identity as a conversion)

Here is the updated example to cover your serialization need, based on the code I've found in your link (with a Serializable base class):

from collections import defaultdict
from dataclasses import dataclass, field
from types import new_class
from typing import (
    Annotated,
    Any,
    AsyncIterable,
    Callable,
    Dict,
    List,
    Mapping,
    Optional,
    TYPE_CHECKING,
    Type,
    TypeVar,
)

import graphql

from apischema import (
    deserialize,
    deserializer,
    schema,
    schema_ref,
    serialize,
    serializer,
)
from apischema.conversions import (
    Conversion,
    dataclass_input_wrapper,
    identity,
    reset_deserializers,
)
from apischema.graphql import graphql_schema
from apischema.metadata import conversion
from apischema.tagged_unions import Tagged, TaggedUnion, get_tagged

# {cls_name: [functions]}
_alternative_constructors: Dict[str, List[Callable]] = defaultdict(list)

if TYPE_CHECKING:
    # Close enough for mypy
    alternative_constructor = staticmethod
else:

    def alternative_constructor(func):
        cls_name = func.__qualname__.rsplit(".", 2)[-2]
        _alternative_constructors[cls_name].append(func)
        return staticmethod(func)


T = TypeVar("T")


# Shortcut
def desc(description: str):
    return schema(description=description)


class TaggedSerializable:
    _deserialization_union: Type[TaggedUnion] = new_class("ScanSpec", (TaggedUnion,))
    _serialization_union: Type[TaggedUnion] = new_class("ScanSpec", (TaggedUnion,))

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        # Retrieved the base class inheriting Serializable
        tagged_base, *_ = [c for c in cls.__mro__ if TaggedSerializable in c.__bases__]
        assert issubclass(tagged_base, TaggedSerializable)
        assert not _, "cannot have multiple base class inheriting directly serializable"
        # If the current class is the base Serializable, do nothing
        if tagged_base == cls:
            return

        serialization = Conversion(
            identity,
            source=cls,
            # Tagged field default serialization (to tagged union) must be bypassed.
            # However, dynamic conversion discards schema_ref, so you must put it back
            # manually, and do the bypass in a sub conversion.
            target=Annotated[cls, schema_ref(cls.__name__)],
            sub_conversions=identity,
        )
        serialization_namespace = {
            "__annotations__": {cls.__name__: Tagged[cls]},
            cls.__name__: Tagged(conversion(serialization=serialization)),
        }
        # Create the serialization tagged union class
        serialization_union = new_class(
            f"Tagged{tagged_base.__name__}Union",
            (tagged_base._serialization_union,),
            exec_body=lambda ns: ns.update(serialization_namespace),
        )
        tagged_base._serialization_union = serialization_union
        # Register the serializer
        serializer(
            Conversion(
                lambda obj: serialization_union(**{obj.__class__.__name__: obj}),
                source=tagged_base,
                target=serialization_union,
            )
        )

        annotations = {cls.__name__: Tagged[cls]}
        deserialization_namespace = {"__annotations__": annotations}
        # Add tagged fields for all its alternative constructors
        for constructor in _alternative_constructors.get(cls.__name__, ()):
            # Deref the constructor function if it is a classmethod/staticmethod
            constructor = constructor.__get__(None, cls)
            # Build the alias of the field
            alias = (
                "".join(map(str.capitalize, constructor.__name__.split("_")))
                + cls.__name__
            )
            # dataclass_input_wrapper uses get_type_hints, but the constructor
            # return type is stringified and the class not defined yet,
            # so it must be assigned manually
            constructor.__annotations__["return"] = cls
            # Wraps the constructor and rename its input class
            wrapper, wrapper_cls = dataclass_input_wrapper(constructor)
            wrapper_cls.__name__ = alias
            # Add constructor tagged field with its conversion
            annotations[alias] = Tagged[cls]
            deserialization_namespace[alias] = Tagged(
                conversion(deserialization=wrapper)
            )
        # Create the deserialization tagged union class
        deserialization_union = new_class(
            f"Tagged{tagged_base.__name__}Union",
            (tagged_base._deserialization_union,),
            exec_body=lambda ns: ns.update(deserialization_namespace),
        )
        tagged_base._deserialization_union = deserialization_union
        # Because deserializers stack, they must be reset before being reassigned
        reset_deserializers(tagged_base)
        # Register the deserializer using get_tagged
        deserializer(
            Conversion(
                lambda obj: get_tagged(obj)[1],
                deserialization_union,
                ScanSpec,
            )
        )

    @classmethod
    def deserialize(cls: Type[T], serialization: Mapping[str, Any]) -> T:
        """Deserialize from a dictionary representation"""
        return deserialize(cls, serialization)

    def serialize(self) -> Mapping[str, Any]:
        """Serialize to a dictionary representation"""
        return serialize(self)


class ScanSpec(TaggedSerializable):
    async def points(self) -> AsyncIterable[float]:
        """Iterate through the points of the scan"""
        raise NotImplementedError


@dataclass
class Line(ScanSpec):
    """A straight line"""

    start: float = field(metadata=desc("The first point"))
    stop: float = field(metadata=desc("The last point"))
    step: float = field(
        default=1, metadata=schema(description="The step between points", exc_min=0)
    )

    async def points(self) -> AsyncIterable[float]:
        point = self.start
        while point <= self.stop:
            yield point
            point += self.step

    @alternative_constructor
    def sized(
        start: Annotated[float, desc("The first point")],
        stop: Annotated[float, desc("The last point")],
        size: Annotated[float, schema(description="Number of points", min=1)],
    ) -> "Line":
        """Alternative constructor with size instead of step"""
        return Line(start=start, stop=stop, step=(stop - start) / (size - 1))


@dataclass
class Concat(ScanSpec):
    left: ScanSpec = field(metadata=desc("First spec to produce"))
    right: ScanSpec = field(metadata=desc("Second spec to produce"))

    async def points(self) -> AsyncIterable[float]:
        async for point in self.left.points():
            yield point
        async for point in self.right.points():
            yield point


def echo(spec: ScanSpec = None) -> Optional[ScanSpec]:
    return spec


scan_spec_schema = graphql_schema(query=[echo])
assert (
    graphql.utilities.print_schema(scan_spec_schema)
    == '''\
type Query {
  echo(spec: TaggedScanSpecUnionInput): TaggedScanSpecUnion
}

type TaggedScanSpecUnion {
  Line: Line
  Concat: Concat
}

type Line {
  """The first point"""
  start: Float!

  """The last point"""
  stop: Float!

  """The step between points"""
  step: Float!
}

type Concat {
  """First spec to produce"""
  left: TaggedScanSpecUnion!

  """Second spec to produce"""
  right: TaggedScanSpecUnion!
}

input TaggedScanSpecUnionInput {
  Line: LineInput
  SizedLine: SizedLineInput
  Concat: ConcatInput
}

input LineInput {
  """The first point"""
  start: Float!

  """The last point"""
  stop: Float!

  """The step between points"""
  step: Float! = 1
}

input SizedLineInput {
  """The first point"""
  start: Float!

  """The last point"""
  stop: Float!

  """Number of points"""
  size: Float!
}

input ConcatInput {
  """First spec to produce"""
  left: TaggedScanSpecUnionInput!

  """Second spec to produce"""
  right: TaggedScanSpecUnionInput!
}
'''
)

query = """\
{
    echo(spec: {
        Concat: {
            left: {
                SizedLine: {
                    start: 0, stop: 12, size: 3
                },
            },
            right: {
                Line: {
                    start: 12, stop: 13
                },
            }
        }
    }) {
        Concat {
            left {
                Line {
                    start stop step
                }
            }
            right {
                Line {
                    start stop step
                }
            }
        }
    }
}
"""

assert graphql.graphql_sync(scan_spec_schema, query).data == {
    "echo": {
        "Concat": {
            "left": {"Line": {"start": 0.0, "stop": 12.0, "step": 6.0}},
            "right": {"Line": {"start": 12.0, "stop": 13.0, "step": 1.0}},
        }
    }
}

Here is another example using (untagged) union for serialization:

from collections import defaultdict
from dataclasses import dataclass, field
from types import new_class
from typing import (
    Annotated,
    Any,
    AsyncIterable,
    Callable,
    Dict,
    List,
    Mapping,
    Optional,
    TYPE_CHECKING,
    Type,
    TypeVar,
    Union,
)

import graphql

from apischema import (
    deserialize,
    deserializer,
    schema,
    schema_ref,
    serialize,
    serialized,
    serializer,
)
from apischema.conversions import (
    Conversion,
    dataclass_input_wrapper,
    identity,
    reset_deserializers,
)
from apischema.graphql import graphql_schema
from apischema.metadata import conversion
from apischema.tagged_unions import Tagged, TaggedUnion, get_tagged

# {cls_name: [functions]}
_alternative_constructors: Dict[str, List[Callable]] = defaultdict(list)

if TYPE_CHECKING:
    # Close enough for mypy
    alternative_constructor = staticmethod
else:

    def alternative_constructor(func):
        cls_name = func.__qualname__.rsplit(".", 2)[-2]
        _alternative_constructors[cls_name].append(func)
        return staticmethod(func)


T = TypeVar("T")


# Shortcut
def desc(description: str):
    return schema(description=description)


class Serializable:
    _deserialization_union: Type[TaggedUnion] = new_class("ScanSpec", (TaggedUnion,))
    _serialization_union: Any = None

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        # Retrieved the base class inheriting Serializable
        tagged_base, *_ = [c for c in cls.__mro__ if Serializable in c.__bases__]
        assert issubclass(tagged_base, Serializable)
        assert not _, "cannot have multiple base class inheriting directly serializable"
        # If the current class is the base Serializable, do nothing
        if tagged_base == cls:
            return

        if tagged_base._serialization_union is None:
            tagged_base._serialization_union = cls
        else:
            tagged_base._serialization_union = Union[
                tagged_base._serialization_union, cls
            ]
            schema_ref("ScanSpecUnion")(tagged_base._serialization_union)
        # Register the serializer
        serializer(
            Conversion(
                identity,
                source=tagged_base,
                target=tagged_base._serialization_union,
                inherited=False,
            )
        )
        # Register an identity serializer to bypass inherited

        annotations = {cls.__name__: Tagged[cls]}
        namespace = {"__annotations__": annotations}
        # Add tagged fields for all its alternative constructors
        for constructor in _alternative_constructors.get(cls.__name__, ()):
            # Deref the constructor function if it is a classmethod/staticmethod
            constructor = constructor.__get__(None, cls)
            # Build the alias of the field
            alias = (
                "".join(map(str.capitalize, constructor.__name__.split("_")))
                + cls.__name__
            )
            # dataclass_input_wrapper uses get_type_hints, but the constructor
            # return type is stringified and the class not defined yet,
            # so it must be assigned manually
            constructor.__annotations__["return"] = cls
            # Wraps the constructor and rename its input class
            wrapper, wrapper_cls = dataclass_input_wrapper(constructor)
            wrapper_cls.__name__ = alias
            # Add constructor tagged field with its conversion
            annotations[alias] = Tagged[cls]
            namespace[alias] = Tagged(conversion(deserialization=wrapper))
        # Create the deserialization tagged union class
        tagged_base._deserialization_union = new_class(
            f"Tagged{tagged_base.__name__}Union",
            (tagged_base._deserialization_union,),
            exec_body=lambda ns: ns.update(namespace),
        )
        # Because deserializers stack, they must be reset before being reassigned
        reset_deserializers(tagged_base)
        # Register the deserializer using get_tagged
        deserializer(
            Conversion(
                lambda obj: get_tagged(obj)[1],
                tagged_base._deserialization_union,
                ScanSpec,
            )
        )

    @classmethod
    def deserialize(cls: Type[T], serialization: Mapping[str, Any]) -> T:
        """Deserialize from a dictionary representation"""
        return deserialize(cls, serialization)

    def serialize(self) -> Mapping[str, Any]:
        """Serialize to a dictionary representation"""
        return serialize(self)


class ScanSpec(Serializable):
    # If you use JSON API too, you can add a serialized property in order to tag objects
    @serialized
    @property
    def typename(self) -> str:
        return self.__class__.__name__

    async def points(self) -> AsyncIterable[float]:
        """Iterate through the points of the scan"""
        raise NotImplementedError


@dataclass
class Line(ScanSpec):
    """A straight line"""

    start: float = field(metadata=desc("The first point"))
    stop: float = field(metadata=desc("The last point"))
    step: float = field(
        default=1, metadata=schema(description="The step between points", exc_min=0)
    )

    async def points(self) -> AsyncIterable[float]:
        point = self.start
        while point <= self.stop:
            yield point
            point += self.step

    @alternative_constructor
    def sized(
        start: Annotated[float, desc("The first point")],
        stop: Annotated[float, desc("The last point")],
        size: Annotated[float, schema(description="Number of points", min=1)],
    ) -> "Line":
        """Alternative constructor with size instead of step"""
        return Line(start=start, stop=stop, step=(stop - start) / (size - 1))


@dataclass
class Concat(ScanSpec):
    left: ScanSpec = field(metadata=desc("First spec to produce"))
    right: ScanSpec = field(metadata=desc("Second spec to produce"))

    async def points(self) -> AsyncIterable[float]:
        async for point in self.left.points():
            yield point
        async for point in self.right.points():
            yield point


def echo(spec: ScanSpec = None) -> Optional[ScanSpec]:
    return spec


scan_spec_schema = graphql_schema(query=[echo])
assert (
    graphql.utilities.print_schema(scan_spec_schema)
    == '''\
type Query {
  echo(spec: TaggedScanSpecUnionInput): ScanSpecUnion
}

union ScanSpecUnion = Line | Concat

type Line {
  """The first point"""
  start: Float!

  """The last point"""
  stop: Float!

  """The step between points"""
  step: Float!
}

type Concat {
  """First spec to produce"""
  left: ScanSpecUnion!

  """Second spec to produce"""
  right: ScanSpecUnion!
}

input TaggedScanSpecUnionInput {
  Line: LineInput
  SizedLine: SizedLineInput
  Concat: ConcatInput
}

input LineInput {
  """The first point"""
  start: Float!

  """The last point"""
  stop: Float!

  """The step between points"""
  step: Float! = 1
}

input SizedLineInput {
  """The first point"""
  start: Float!

  """The last point"""
  stop: Float!

  """Number of points"""
  size: Float!
}

input ConcatInput {
  """First spec to produce"""
  left: TaggedScanSpecUnionInput!

  """Second spec to produce"""
  right: TaggedScanSpecUnionInput!
}
'''
)

query = """\
{
    echo(spec: {
        Concat: {
            left: {
                SizedLine: {
                    start: 0, stop: 12, size: 3
                },
            },
            right: {
                Line: {
                    start: 12, stop: 13
                },
            }
        }
    }) {
        __typename
        ... on Concat {
            left {
                __typename
                ... on Line {
                    start stop step
                }
            }
            right {
                __typename
                ... on Line {
                    start stop step
                }
            }
        }
    }
}
"""

assert graphql.graphql_sync(scan_spec_schema, query).data == {
    "echo": {
        "__typename": "Concat",
        "left": {"__typename": "Line", "start": 0.0, "stop": 12.0, "step": 6.0},
        "right": {"__typename": "Line", "start": 12.0, "stop": 13.0, "step": 1.0},
    }
}

[thomascobb]

Thanks, that fixes it. I've modified it slightly and put the changes in a PR for comment:
https://github.com/dls-controls/scanspec/pull/9/files

There's only one part I don't understand, why does the Conversion(identity) need sub_converters=identity? I stepped through the serialize code, and can see that it is called a number of times:

• serialize(line, converters=Conversion(to_tagged_union))
• serialize(tagged_union(Line=line), converters=None)
• serialize(line, converters=Conversion(identity))
• serialize(line, converters=None)

I guess that last serialize is why sub_converters=identity is needed, but I assumed that converters=Conversion(identity) would do that?

[wyfo]

You're right to highlight this point, it's in fact quite complex.
Actually, if GraphQL were not involved (if you only serialize or serialization_schema), Conversion(identity) would be enough to turn off the registered serializer.
However, the issue with GraphQL schema generation is resumed in my comment:

serialization = Conversion(
            identity,
            source=cls,
            # Tagged field default serialization (to tagged union) must be bypassed.
            # However, dynamic conversion discards schema_ref, so you must put it back
            # manually, and do the bypass in a sub conversion.
            target=Annotated[cls, schema_ref(cls.__name__)],
            sub_conversions=identity,
        )

Let me try to explain with more details:
Each type present in the GraphQL schema must be named. By default, apischema takes the name of the class mapped to the type, even if this type has registered conversion. However, in case of dynamic conversion, the name of the class cannot be used (because it's reserved for the type without dynamic conversion), so name of the conversion target is used instead, but not in case of registered conversion bypass, because it would conflict with the type without dynamic conversion.
If you replace the line by serialization = identity, just to do the bypass, you will get an error saying Missing ref for type <class '__main__.Line'>: you have no name for the type mapped from the Line class (the one resulting from the conversion bypass).
That's why you need to a schema_ref manually, which is done with target=Annotated[cls, schema_ref(cls.__name__)],, and because this conversion does nothing else, it use the identity function (but it could also be lambda x: x here). Then, because the conversion to tagged union must be bypassed, you need to do it in a sub conversion with sub_conversions=identity (lambda x: x would not work, because apischema use identity as a marker for conversion bypass). You have then a conversion to an annotated type (whose schema_ref will be used as the type name in the schema), which is converted to the raw type by bypassing the registered conversion.

All of that because you need to turn off the conversion of the tagged union member to avoid looping because it would be other converted to the tagged union.

[thomascobb]

Thanks, that explains it. I'm a lot happier with this solution registering both serializers and deserializers, it looks cleaner to me.

[thomascobb]

Actually, in my opinion, the best thing would be to not use either of them but extract the descriptions (and why not deprecation) directly from the docstring of classes/functions, because I think it's the right place for field/parameter descriptions.

Personally, I'm not a big fan of extracting information from docstrings as there seem to be many standards. I use sphinx, so we already have the :param: syntax as well as the numpy and google style docstrings supported by sphinx.ext.napoleon

I decided to go the other way, and extract the apischema descriptions for the docs. I think you've seen some of the code already as it was extracted from scanspec. Here's the repo:
https://github.com/dls-controls/sphinx-apischema/