Proper support for functional discriminators

[samuelcolvin]

Initial Checks

• I have searched Google & GitHub for similar requests and couldn't find anything
• I have read and followed the docs and still think this feature is missing

Description

this is supported in pydantic-core, see here but doesn't seem to be well supported in Pydantic V2.

I propose we add a new type:

@dataclass
class Discriminator:
    discriminator: Callable[[Any], str] | str
    _: KW_ONLY
    # defaults to `{choice.__name__: choice for choice in choices}`
    choice_lookup: Callable[tuple[type[Any], ...], dict[str, type[Any]]] | None = None

Obviously Discriminator should behave the same as Field(discriminator='...') when it's used with a string.

Which could be used like this:

from dataclasses import dataclass
from typing import Annotated

from pydantic import TypeAdapter, Discriminator


@dataclass
class City:
    name: str


@dataclass
class Country:
    name: str


def discriminator_function(obj: Any) -> str:
    if obj.name == 'London':
        return 'City'
    else:
        return 'Country'


ts = TypeAdapter(Annotated[City | Country, Discriminator(discriminator_function)])

Related to #7391 but deserved a better explanation.

Affected Components

• Compatibility between releases
• Data validation/parsing
• Data serialization - .model_dump() and .model_dump_json()
• JSON Schema
• Dataclasses
• Model Config
• Field Types - adding or changing a particular data type
• Function validation decorator
• Generic Models
• Other Model behaviour - model_construct(), pickling, private attributes, ORM mode
• Plugins and integration with other tools - mypy, FastAPI, python-devtools, Hypothesis, VS Code, PyCharm, etc.

[karlicoss]

Ah sorry if I was a bit unclear in #7391
Basically I'd like a way for any union to work (including any recursive union types, i.e. if it's a union field in a dataclass etc) without explicitly adding discriminators for them (I'm not sure if it somehow goes against pydantic philosophy though)

In this particular case I can see how your suggestion would work for serialization -- however not sure it's gonna work for deserializing (ts.validate_python call)? Perhaps you could add a suggested output of ts.dump_python when using a functional discriminator -- I think it might be a bit more clear then?

[sydney-runkle]

Side note -- looks like David has already made significant progress in this department: #6915.

Will touch base with the team on this tmrw morning.

[JacobHayes]

Would this (or something else) support dynamically registering members that are not know when defining the Union (or equivalent)? Of course, all members would still have to be known by the time of (de)serialization (or perhaps dynamically registered by discriminator_function if possible).

Specifically, I'm unsure whether choices/choice_lookup could safely be extended after used in a type hint?

Happy to open/move to a separate issue if this isn't the right place. (edit: #7366 might be the better issue, but is closed and points here.)

---

For example, I've worked on 3 unrelated projects where a base class defines a couple fields but is then "dynamically" subclassed with more fields. In two of the projects, external users are creating the subclasses outside the library, hence can't be know ahead of time (but we can implement a plugin/registration mechanism to ensure they are know by the time we try (de)serialization). The other project uses separate modules to avoid circular imports, so maintaining the Union definition each time we add a module/subclass is extra challenging.

The reason we'd like a discriminated union here, instead of just hinting with the base class, is because of these additional fields defined on the subclasses. If a wrapping model is just hinted with the base class, the extra fields are discarded.

Perhaps just updating Discriminator.choices dynamically and using the base class as the type hint could work? I'm guessing just having the base class as the hint might not play too well with generated JSON Schema doc (unless that is predicated on the choices?). Eg:

class Base(BaseModel):
    kind: str

base_discriminator = Discriminator(...)

class Container(BaseModel):
    data: Annotated[Base, base_discriminator]

... # elsewhere

class A(Base):
    kind: Literal["A"] = "A"
    a: int

base_discriminator.choices["A"] = A

class B(Base):
    kind: Literal["B"] = "B"
    b: int

base_discriminator.choices["B"] = B

(in reality, probably made easier with __init_subclass__ registration, etc).

[sydney-runkle]

Ta da! Closed by #7983 🥳