Narrowing types to Literal using in syntax

[tusharsadhwani]

Take this example code:

x = "a"
reveal_type(x)
assert x in ("a", "b", "c")
reveal_type(x)

y = 20
reveal_type(y)
assert y == 10 or y == 20 or y == 50
reveal_type(y)

The current output of this is:

$ mypy x.py 
x.py:2: note: Revealed type is "builtins.str"
x.py:4: note: Revealed type is "builtins.str"
x.py:7: note: Revealed type is "builtins.int"
x.py:9: note: Revealed type is "builtins.int"

If mypy were to support type narrowing using chained or statements and in statements, the output would look like:

$ mypy x.py 
x.py:4: note: Revealed type is "builtins.str"
x.py:6: note: Revealed type is "Union[Literal['a'], Literal['b'], Literal['c']]"
x.py:9: note: Revealed type is "builtins.int"
x.py:11: note: Revealed type is "Union[Literal[10], Literal[20], Literal[50]]"

I think doing simple narrowing here would be relatively easy and very useful.

[tmke8]

Related to (but maybe not quite a duplicate of) #9718.

[tusharsadhwani]

It's a superset of #9718. I'll be looking into implementing this then, starting with the literal narrowing case.

[AlexWaygood]

A superset of #9718, and a subset of #10977 :)

[bluenote10]

This issue seems to address narrowing both for in but also for ==. Should these two features be tracked independently (if their implementations would be rather decoupled) or is tracking them together preferred (if implementing one would likely solve the other)?

I couldn't find an issue specifically for narrowing literals in == equality comparisons, and I would have expected this to be implemented first, because it even feels more basic than the in operator:

from typing import Literal

def takes_foo(s: Literal["foo"]) -> None:
    ...

def f(s: str) -> None:
    if s == "foo":
        takes_foo(s)

error: Argument 1 to "takes_foo" has incompatible type "str"; expected "Literal['foo']"  [arg-type]

Pyright seems to understand this kind of narrowing well.

[asottile-sentry]

I opted for a clunky workaround using TypeIs though it would be nice if this were automatic

MyLiteral = Literal["a", "b"]

def is_my_literal(s: str) -> TypeIs[MyLiteral]:
    return s in ("a", "b")

[JelleZijlstra]

This is similar to #3229 though that one doesn't ask for narrowing to Literal, possibly because the issue predates literal types.