Skip to content

Conversation

@YuriiMotov
Copy link
Member

This PR updates type hints for discriminator parameter of Field to be in line with Pydantic's version:

    discriminator: str | types.Discriminator | None = _Unset,

Also, added tests for callable discriminator.


The following code works on master:

from collections.abc import Hashable
from typing import Annotated, Any, Literal, Union

from pydantic import Discriminator, Tag
from sqlmodel import Field, SQLModel

class Pie(SQLModel):
    time_to_cook: int
    num_ingredients: int

class ApplePie(Pie):
    fruit: Literal["apple"] = "apple"

class PumpkinPie(Pie):
    filling: Literal["pumpkin"] = "pumpkin"

def get_discriminator_value(v: Any) -> Hashable:
    if isinstance(v, dict):
        return v.get("fruit", v.get("filling"))
    return getattr(v, "fruit", getattr(v, "filling", None))

    class ThanksgivingDinner(SQLModel):
        dessert: Union[
            Annotated[ApplePie, Tag("apple")],
            Annotated[PumpkinPie, Tag("pumpkin")],
        ] = Field(
            discriminator=Discriminator(get_discriminator_value),
        )

apple_variation = ThanksgivingDinner.model_validate(
    {"dessert": {"fruit": "apple", "time_to_cook": 60, "num_ingredients": 8}}
)

pumpkin_variation = ThanksgivingDinner.model_validate(
    {"dessert": {"filling": "pumpkin", "time_to_cook": 40, "num_ingredients": 6}}
)

print(type(apple_variation.dessert))
# <class '__main__.ApplePie'>

print(type(pumpkin_variation.dessert))
# <class '__main__.PumpkinPie'>

.. but mypy argues with:

error: No overload variant of "Field" matches argument types "EllipsisType", "Discriminator"  [call-overload]

This PR fixes this

@YuriiMotov YuriiMotov added the feature New feature or request label Jan 29, 2026
@YuriiMotov YuriiMotov changed the title Allow Discriminator for discriminator in Field 🏷️ Allow Discriminator for discriminator in Field Jan 29, 2026
@YuriiMotov YuriiMotov marked this pull request as ready for review January 29, 2026 09:29
@YuriiMotov YuriiMotov force-pushed the discriminator-type-hints branch from 80d2b66 to 80f8986 Compare January 29, 2026 09:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants