Narrowing with "tags" on unions (of TypedDicts or normal classes) doesn't work with the match statement

[tmke8]

Bug Report

Mypy has a feature where a union of TypedDicts can be narrowed down based on the value of entries: https://mypy.readthedocs.io/en/stable/literal_types.html#tagged-unions

But this doesn't seem to work with pattern matching.

To Reproduce

from typing import Literal, TypedDict

class A(TypedDict):
    tag: Literal["a"]
    name: str
    
class B(TypedDict):
    tag: Literal["b"]
    num: int
    
def f(d: A | B) -> None:
    if d["tag"] == "a":
        print(d["name"])
    elif d["tag"] == "b":
        print(d["num"])

def g(d: A | B) -> None:
    match d["tag"]:
        case "a":
            print(d["name"])  # E: TypedDict "B" has no key "name"
        case "b":
            print(d["num"])  # E: TypedDict "A" has no key "num"

https://mypy-play.net/?mypy=latest&python=3.11&gist=22cd1c2a185a5182b98e12c1e5dd33e2

Expected Behavior

The narrowing in function g should work the same as in function f.

Actual Behavior

main.py:20: error: TypedDict "B" has no key "name"  [typeddict-item]
main.py:22: error: TypedDict "A" has no key "num"  [typeddict-item]
Found 2 errors in 1 file (checked 1 source file)

The dictionaries are not correctly narrowed in the match branches.

Your Environment

See mypy-play link above.

[u7601139]

@AlexWaygood @tmke8 @JelleZijlstra Hi, may I have a try on this issue?

[AlexWaygood]

@AlexWaygood @tmke8 @JelleZijlstra Hi, may I have a try on this issue?

Yes, feel free!

[richardxia]

I'm not sure if I should file this as a separate bug ticket, but I noticed the exact same behavior for normal classes as well, not just TypedDicts. I tweaked the example code from the issue description to demonstrate this, showing that the narrowing works fine for if statements but not match statements for regular classes.

from typing import Literal

class A:
    tag: Literal["a"]
    name: str
    
class B:
    tag: Literal["b"]
    num: int
    
def f(d: A | B) -> None:
    if d.tag == "a":
        print(d.name)
    elif d.tag == "b":
        print(d.num)

def g(d: A | B) -> None:
    match d.tag:
        case "a":
            print(d.name)  # E: Item "B" of "A | B" has no attribute "name"
        case "b":
            print(d.num)  # E: Item "A" of "A | B" has no attribute "num"

Output:

main.py:20: error: Item "B" of "A | B" has no attribute "name"  [union-attr]
main.py:22: error: Item "A" of "A | B" has no attribute "num"  [union-attr]
Found 2 errors in 1 file (checked 1 source file)

https://mypy-play.net/?mypy=latest&python=3.12&gist=1b66ee945421120ea4215eaf2ae072a1

[tmke8]

I can change the title of the issue to reflect the broader scope.

[rijenkii]

FYI -- this has just been fixed in pyright 1.1.342:

rijenkii@rijenkiipc /tmp> cat test.py 
from typing import Literal

class A:
    tag: Literal["a"]
    name: str

class B:
    tag: Literal["b"]
    num: int

def f(d: A | B) -> None:
    if d.tag == "a":
        reveal_type(d)
    elif d.tag == "b":
        reveal_type(d)

def g(d: A | B) -> None:
    match d.tag:
        case "a":
            reveal_type(d)
        case "b":
            reveal_type(d)

rijenkii@rijenkiipc /tmp> PYRIGHT_PYTHON_FORCE_VERSION=latest pyright --version
pyright 1.1.342

rijenkii@rijenkiipc /tmp> PYRIGHT_PYTHON_FORCE_VERSION=latest pyright test.py
/tmp/test.py
  /tmp/test.py:13:21 - information: Type of "d" is "A"
  /tmp/test.py:15:21 - information: Type of "d" is "B"
  /tmp/test.py:20:25 - information: Type of "d" is "A"
  /tmp/test.py:22:25 - information: Type of "d" is "B"
0 errors, 0 warnings, 4 informations