Nested tagged unions?

[uar316025]

Question

Hi! Can i do something like that? (not sure how to name it) Here example with pydantiic (v2)

First tagged union by type field Then second by sub_type field

from typing import Literal, Annotated

import pydantic

class Device(pydantic.BaseModel):
    type: int
    sub_type: int
    name: str

class BlenderM1(Device):
    type: Literal[0x01] = 0x01
    sub_type: Literal[0x01] = 0x01

class BlenderM2(BlenderM1):
    sub_type: Literal[0x02] = 0x02

class OvenM1(Device):
    type: Literal[0x02] = 0x02
    sub_type: Literal[0x01] = 0x01

class OvenM2(OvenM1):
    sub_type: Literal[0x02] = 0x02

AnyDevice = Annotated[
    Annotated[BlenderM1 | BlenderM2, pydantic.Field(discriminator='sub_type')] |
    Annotated[OvenM1 | OvenM2, pydantic.Field(discriminator='sub_type')],
    pydantic.Field(discriminator='type')
]
print(pydantic.TypeAdapter(list[AnyDevice]).validate_python([
    {'type': 0x01, 'sub_type': 0x01, 'name': 'BM1'},
    {'type': 0x01, 'sub_type': 0x02, 'name': 'BM2'},
    {'type': 0x02, 'sub_type': 0x01, 'name': 'OM1'},
    {'type': 0x02, 'sub_type': 0x02, 'name': 'OM2'},
]))

which prints

[
  BlenderM1(type=1, sub_type=1, name='BM1'),
  BlenderM2(type=1, sub_type=2, name='BM2'),
  OvenM1(type=2, sub_type=1, name='OM1'),
  OvenM2(type=2, sub_type=2, name='OM2')
]

Also i can't change the devices structure because it is defined externally.

[jwhitaker-gridcog]

yeah.. pydantic gets this one right I think - in that the tag field is a property of the union, not of each struct. With that property, this is straightforward to represent. OpenAPI and an upcoming revision of JSONSchema also have this as a property of the union.

Msgspec's decision to have the tag field be a property of each struct instead might make this unreasonably difficult. Maybe that's impetus to re-examine https://github.com/jcrist/msgspec/issues/338 ? Putting this on the union instead (with an Annotated like pydantic, or by some other method, possibly even automatic?) could maybe open the door to allowing unions of TypedDicts, dataclasses, etc as well.