Constructing tagged-union typeddicts with Any

[bijij]

I have a question about type-narrowing and tagged-union typeddicts, specifically with Pyright I suppose narrowing earlier than I would like,

In a project I have a Payload sent to (and from) an API with multiple keys omitted or included depending on the type key value.
e.g. (in reality there are about 10 different types at this point)

class _BaseFooDict(TypedDict):
    name: str


class _FooDictA(_BaseFooDict):
    type: Literal[1]


class _FooDictBOptional(_BaseFooDict, total=False):
    description: Optional[str]


class _FooDictB(_FooDictBOptional):
    type: Literal[2]


# ... etc


FooDict = Union[_FooDictA, _FooDictB]  # ... etc

I use a tagged-union of these typeddicts as the return parameter of a function which generates this from an instance of a class.

mypy will spit out an error on the line defining foo_dict saying the Type is ambiguous, which, is true (not sure i agree with erroring there but so-be it, the type revealed below however will be corectly interprited as _FooDictB.

Something like:

class FooType(Enum):
    A = 1
    B = 2
    ...  # etc


class Foo:
    name: str
    type: FooType

    # required for type B
    description: Optional[str]

    def to_dict(self) -> FooDict:
        foo_dict: FooDict = {
            "name": self.name,
            "type": self.type.value,
        }

        if self.description is not None:
            assert foo_dict["type"] == 2
            # reveal_type(foo_dict) 
            foo_dict["description"] = self.description

        return foo_dict

My issue is, in the if self.description is not None block, after the assert the revealed type will be Never because pyright has I'm assuming from the self.type.value being Any lazily narrowed the tagged-union to the first option _FooDictA rather than keeeping the Union intact.

I'm opening this discussion to ask if anyone has a suggestion on how to, effectivly keep the union intact in this situation, I could use cast(FooDict, ...) to reset the narrowing just after the fact but that sort of feels clunky to me, other option I've got is just to ignore the issue since it's in Never land the type-checker won't care anyway, but again doesn't feel right to me.

Thoughts?

[erictraut]

There are a few problems with your approach.

Pyright special-cases the value field of an enum literal type, so the type of FooType.A.value is known (evaluated to Literal[1]), but you're accessing the value field from an unknown (dynamic) enum member, so the value type is Any. One way around this is to avoid using an enum for this use case. Switch to a type alias that includes literals for each value.

FooType = Literal[1, 2]

Second, when you assign a dict expression to a TypedDict (as in foo_dict: FooDict = {...}), the dict expression on the right must be compatible with the TypedDict annotation on the left. The type of the dict expression in your case is:

{
    "name": <value type is str>,
    "type": <value type is Literal[1] | Literal[2]>
}

There is no TypedDict that supports a type field with a union type, so the assignment fails.

So here's one solution that eliminates all type errors:

FooType = Literal[1, 2]

class Foo:
    name: str
    type: FooType

    # required for type B
    description: Optional[str]

    def to_dict(self) -> FooDict:
        foo_dict: FooDict
        if self.type == 1:
            foo_dict = {"name": self.name, "type": 1}
        elif self.type == 2:
            foo_dict = {"name": self.name, "type": 2, "description": self.description}
        else:
            raise ValueError('Unexpected type')

        return foo_dict

This approach would also work with enums:

class FooType(Enum):
    A = 1
    B = 2
    ...  # etc

class Foo:
    name: str
    type: FooType

    # required for type B
    description: Optional[str]

    def to_dict(self) -> FooDict:
        foo_dict: FooDict
        if self.type == FooType.A:
            foo_dict = {"name": self.name, "type": 1}
        elif self.type == FooType.B:
            foo_dict = {"name": self.name, "type": 2, "description": self.description}
        else:
            raise ValueError('Unexpected type')

        return foo_dict

I don't fully understand what you're trying to do here, but it seems odd that the class Foo would have knowledge about all of the different types of TypedDicts. Would it make sense to break this into subclasses like A and B and have each class know how to construct its corresponding TypedDict type — i.e. the to_dict method implemented in class A would produce _FooDictA, and class B would produce _FooDictB?