Skip to content

Typing breaks on Pydantic discriminated unions using Dagster Config objects as base  #25694

Open
@carterinvres

Description

@carterinvres

What's the issue?

When using a typical Pydantic discriminated union over Dagster Config classes via the
mytype = Annotated[Union[class1, class2], Field(discriminator="method"),
the resulting type does not appear to have any typehinting information and the indicated type is Unknown.

What did you expect to happen?

If a Pydantic discriminated union is created across multiple BaseModel types using the
mytype = Annotated[Union[class1, class2], Field(discriminator="method"),
then objects of that type have the typehint mytype. Attempting to call a function on an object of that type will also raise typing errors if one of the variants in the union does not have that function implemented/if return types don't match expected return type.

I would expect that union types created using the Annotated[..., discriminator=] syntax with Dagster Config objects would behave the same way.

How to reproduce?

from dagster import Config, ConfigurableResource
from typing import Literal, Annotated, Union
from pydantic import Field, BaseModel


class Empty(BaseModel):
    method: Literal["EMPTY"] = "EMPTY"


class NotEmpty(BaseModel):
    method: Literal["NOT_EMPTY"] = "NOT_EMPTY"


ResourceBase = Annotated[Union[Empty, NotEmpty], Field(discriminator="method")]


class MyBaseModelType(BaseModel):
    base: ResourceBase

    def return_base(self):
        return self.base


class EmptyConfig(Config):
    method: Literal["EMPTY"] = "EMPTY"


class NotEmptyConfig(Config):
    method: Literal["NOT_EMPTY"] = "NOT_EMPTY"


ResourceBaseConfig = Annotated[Union[EmptyConfig, NotEmptyConfig], Field(discriminator="method")]


class MyConfigurableResource(ConfigurableResource):
    base: ResourceBaseConfig

    def return_base(self):
        return self.base

The signature for MyBaseModelType.return_base is ResourceBase, while the signature for MyConfigurableResource.return_base function is Any

Dagster version

version 1.8.5

Deployment type

None

Deployment details

Deployment shouldn't be relevant - the code itself works fine, just the typehints are broken in IDE.

Additional information

No response

Message from the maintainers

Impacted by this issue? Give it a 👍! We factor engagement into prioritization.
By submitting this issue, you agree to follow Dagster's Code of Conduct.

Metadata

Metadata

Assignees

No one assigned

    Labels

    area: configRelated to Configurationtype: bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions