Make pyright recognise narrowing down a type using a in-place set

[DanielNoord]

Is your feature request related to a problem? Please describe.

from typing import Literal

def func(param: int) -> Literal[1, 3]:
    assert param in (1, 3)
    reveal_type(param) # Reveals Literal[1,3]
    return param

def func_two(param: int) -> Literal[1, 3]:
    assert param in {1, 3}
    reveal_type(param) # Reveals int
    return param

Describe the solution you'd like
Please correct me if I'm wrong, but shouldm't the type also be narrowed down to Literal[1,3] in the second function?

Additional context

def func_three(param: int) -> Literal[1]:
    assert param in (1)
    reveal_type(param) # Reveals int
    return param

Perhaps this is trickier because it is a tuple with a single value and there can be some edge-cases I can't immediately think of, but shouldn't this also narrow down to Literal[1] like func does?

[erictraut]

This won't work (at least not without some really fragile special-case logic) because, unlike tuples, the inferred type of a set doesn't preserve literals.

x1 = (1, 3)
reveal_type(x1) # tuple[Literal[1], Literal[3]]
x2 = {1, 3}
reveal_type(x2) # set[int]

If you want to use this type narrowing pattern in your code, you'd need to switch from sets to tuples.

[DanielNoord]

Is that by design? Or can a proposal be made to retain the literal types within a set?

The second example is also non-relevant I assume?

[erictraut]

That's by design. You really wouldn't want a type checker to infer the type of {1, 3} to be set[Literal[1] | Literal[3]] because that would prevent you from adding anything to the set other than a 1 or 3. Tuples are immutable and the values of individual entries within a tuple are tracked separately, whereas set has a single type parameter that specifies types of all values within the set.

If you really want to use a set here for some reason, you could do the following:

def func_two(param: int) -> Literal[1, 3]:
    x: set[Literal[1, 3]] = {1, 3}
    assert param in x
    reveal_type(param)
    return param
``

[DanielNoord]

Hmm okay. I ran into this since pylint suggests using a set for membership checks because of performance differences. See pylint-dev/pylint#4841 and pylint-dev/pylint#4776 for discussion about this.

I can see how sets defined anywhere else shouldn't be inferred to set[Literal[1] | Literal[3]] but when the set is defined in place (and thus cast to frozenset as discussed in issue 4776) it might make sense to do so? I'm not fully aware of Python's internals as some of the people that commented in those discussions, but to me it could make sense to infer the literals in a frozenset.