Strange behavior using Tagged Union

[msdrigg]

Bug Report

When using tagged unions, mypy infers correctly with with instance["tag"],
but fails with instance.get("tag")

To Reproduce

1. Write the following code into tagged_union.py:

from typing import Literal, TypedDict, Union


class TaggedTypeA(TypedDict):
    tag: Literal["A"]
    is_a: bool


class TaggedTypeB(TypedDict):
    tag: Literal["B"]


TaggedUnion = Union[TaggedTypeA, TaggedTypeB]

instance_a: TaggedUnion = {"tag": "A", "is_a": True}

if instance_a["tag"] == "A":
    print(f"Am I type a? {instance_a['is_a']}")

if instance_a.get("tag") == "A":
    print(f"Am I type a this time? {instance_a['is_a']}")

2. Run mypy tagged_union.py
3. Get only one error

tagged_union.py:21: error: TypedDict "TaggedTypeB" has no key "is_a"

Error occurs in second if statement

Expected Behavior

I can see that mypy correctly infers the type of instance_a when I use instance_a["tag"] == "A" (the first if expression).

Therefore, I would expect the same behavior when using instance_a.get("tag") == "A" (the second if clause).

Actual Behavior
mypy correctly infers the type of instance_a when I use instance_a["tag"] == "A" (the first if expression).

mypy fails to infer the type when I use instance_a.get("tag") == "A".

In fact, mypy doesn't only not infer the correct type, it infers the opposite type TaggedTypeB which does not make sense at all

Your Environment

• Mypy version used: 0.902
• Mypy command-line flags: None
• Mypy configuration options from mypy.ini (and other config files): None
• Python version used: 3.8.10
• Operating system and version: Ubuntu 21.04

[ethframe]

It's not that mypy narrows the type incorrectly - it's just an imprecise error message:

...

# No narrowing occurs here since mypy does not understand this pattern
if instance_a.get("tag") == "A":
    reveal_type(instance_a)  # Revealed type is "Union[TaggedTypeA, TaggedTypeB]"
    print(f"Am I type a this time? {instance_a['is_a']}")  # TypedDict "TaggedTypeB" has no key "is_a"

Compare this to accessing an attribute of union type:

from typing import Union


class A:
    a: int
    b: bool


class B:
    a: int


ab: Union[A, B] = A()
ab.b  # Item "B" of "Union[A, B]" has no attribute "b"

[msdrigg]

Thank you for the explanation. This error message makes sense given the context.

Two follow up questions before I close this:

1. Should the error message be clarified to clearly indicate that the error is due to Union including B, not that the type has been transformed into B? I dont know if this change should be made to account for a small subset of errors.
2. Should we be checking for TypedDict.get(Literal) to convert unions? The single argument .get(Literal) case is equivalent to the indexing operator for the purposes of automatic casting. Should this case be implemented?

[levrik]

Since the upgrade to mypy 0.900 I'm also seeing this behavior while before the narrowing was working correctly.
I'm accessing attributes on a class. Pyright (alternative type checker by Microsoft) for example gets it right as mypy before.

[msdrigg]

I would be willing to work on a pull request on this issue if someone could point me to where in the codebase is responsible for handling the transformation between union and specific types.

[hauntsaninja]

@levrik if you're seeing a regression, please open another issue with your code. Previous mypy versions complain about the examples in this issue, so seems separate from what you're describing.

As for implementing narrowing for tagged unions using .get maybe try looking at

mypy/mypy/checker.py

Line 4348 in 3319826

def refine_parent_types(self,

[levrik]

@hauntsaninja I've opened #10764 for the issue I'm seeing.