Tagged unions don't work with Enums

[hackaugusto]

Bug Report

With the example below, the second if doesn't properly narrow the type, it works fine if the Literal['b'] instead of the enum version.

To Reproduce

from queue import Queue
from typing import Literal, TypedDict, Union
from enum import Enum, unique

@unique
class Tag(Enum):
 a = 'a'
 b = 'b'

class OpA(TypedDict):
 tag: Literal[Tag.a]
 a: str

class OpB(TypedDict):
 tag: Literal[Tag.b]
 b: str

Ops = Union[OpA, OpB]

def do(op: Ops):
 if op['tag'] == Tag.a:
 reveal_type(op)
 if op['tag'] == Tag.b:
 reveal_type(op)

Expected Behavior

 if op['tag'] == Tag.b:
 reveal_type(op)

Should report:

t.py:24: note: Revealed type is "TypedDict('t.OpB', {'tag': Literal[t.Tag.b], 'b': builtins.str})"

Actual Behavior

Reported type is:

t.py:24: note: Revealed type is "Union[TypedDict('t.OpA', {'tag': Literal[t.Tag.a], 'a': builtins.str}), TypedDict('t.OpB', {'tag': Literal[t.Tag.b], 'b': builtins.str})]"

Your Environment

• Mypy version(s) used: mypy 0.910, mypy 0.941, mypy 0.950+dev.fd8a15e1abdd2bb4f75c3a6319ebfba233e1fe8e (compiled: no)
• Mypy command-line flags: mypy t.py
• Mypy configuration options from mypy.ini (and other config files): nothing
• Python version used: Python 3.10.2
• Operating system and version: Fedora Linux 35

[AlexWaygood]

Although narrowing using == is completely valid, and mypy is indeed incorrect here, the more idiomatic way to compare enum values is using is rather than ==. If you use is instead of ==, mypy understands your code flow as you'd expect (mypy playground link here).

You'd just rewrite your do function like so:

def do(op: Ops):
 if op['tag'] is Tag.a:
 reveal_type(op)
 if op['tag'] is Tag.b:
 reveal_type(op)

[AlexWaygood]

With regard to mypy not understanding enum narrowing using ==, there's already an open PR by @hauntsaninja to fix that: #11521

[AlexWaygood]

Closing as a duplicate of #10915