Narrowing unions with in/__contains__

[sfoster1]

It seems like you should be able to narrow a union of literals to a subset, or a union of unions of literals to just one of the unions, by checking if a value with such a union is in a container of one of the subset/sub unions (see the code example for something hopefully more coherent).

This is really useful for, say, a deserialized JSON object that has different subsets of keys that map to differently-shaped objects, and you want to pass them to functions that take the appropriate parameters.

• Are you reporting a bug, or opening a feature request?

Really could be either - this is behavior I'm surprised by, but it could be not implemented rather than working incorrectly

• Minimal repro showing the issue:

Playground gist link: https://mypy-play.net/?mypy=latest&python=3.8&gist=18262b0bbd541a22d4f48c5bb628f00e
Code:

from typing import Union, Dict, List, Literal

firstset = Union[Literal['a'], Literal['b'], Literal['c']]
secondset = Union[Literal['d'], Literal['e'], Literal['f']]

firstdict: Dict[firstset, int] = {'a': 2}
seconddict: Dict[secondset, float] = {'d': 3.0}

firstlist: List[firstset] = ['a', 'b', 'c']
secondlist: List[secondset] = ['d', 'e', 'f']

def printit(key: Union[Literal['a'], Literal['b'], Literal['c'],
                       Literal['d'], Literal['e'], Literal['f']]):
    if key in firstdict:
        reveal_type(key) # Union[Literal['a'], Literal['b'], Literal['c'], Literal['d'], Literal['e'], Literal['f']]'
    if key in seconddict:
        reveal_type(key) # Union[Literal['a'], Literal['b'], Literal['c'], Literal['d'], Literal['e'], Literal['f']]'

    if key in firstlist:
        reveal_type(key) # Union[Literal['a'], Literal['b'], Literal['c'], Literal['d'], Literal['e'], Literal['f']]'
    if key in secondlist:
        reveal_type(key) # Union[Literal['a'], Literal['b'], Literal['c'], Literal['d'], Literal['e'], Literal['f']]'

    if key == 'a' or key == 'b' or key == 'c':  # this one works
        reveal_type(key)  # 'Union[Literal['a'], Literal['b'], Literal['c']]' 

• What is the actual behavior/output?

test.py:16: note: Revealed type is 'Union[Literal['a'], Literal['b'], Literal['c'], Literal['d'], Literal['e'], Literal['f']]'
test.py:18: note: Revealed type is 'Union[Literal['a'], Literal['b'], Literal['c'], Literal['d'], Literal['e'], Literal['f']]'
test.py:21: note: Revealed type is 'Union[Literal['a'], Literal['b'], Literal['c'], Literal['d'], Literal['e'], Literal['f']]'
test.py:23: note: Revealed type is 'Union[Literal['a'], Literal['b'], Literal['c'], Literal['d'], Literal['e'], Literal['f']]'
test.py:26: note: Revealed type is 'Union[Literal['a'], Literal['b'], Literal['c']]'

• What is the behavior/output you expect?

test.py:16: note: Revealed type is 'Union[Literal['a'], Literal['b'], Literal['c']]
test.py:18: note: Revealed type is 'Union[Literal['d'], Literal['e'], Literal['f']]
test.py:21: note: Revealed type is 'Union[Literal['a'], Literal['b'], Literal['c']]
test.py:23: note: Revealed type is 'Union[Literal['d'], Literal['e'], Literal['f']]
test.py:26: note: Revealed type is 'Union[Literal['a'], Literal['b'], Literal['c']]'

• What are the versions of mypy and Python you are using?
  Do you see the same issue after installing mypy from Git master?
  Happens on 3.7 and 3.8, mypy 0.770 and whatever playground has as "latest"
• What are the mypy flags you are using? (For example --strict-optional)
  Defaults

(FYI if you were reading this shortly after I posted it and were annoyed by me editing it a lot - I'm really sorry, I accidentally posted it early and needed to actually fill it out)

[emmatyping]

I think this seems reasonable to support, but I'm not sure how hard it would be. Perhaps @Michael0x2a has a better idea.

[antonagestam]

This shares some aspects with my proposal to make immutable containers work like literal unions: #8689

[sfoster1]

@ethanhs Opened a PR with something that at least begins to implement the issue - interested in your (or @Michael0x2a 's, or really anybody's) thoughts.

@antonagestam I think this would be quite similar to your feature request but would at least at the moment require extra annotation in your code, like

ALLOWED_METHODS: Final[FrozenSet[Union[Literal['POST']...]]] = frozenset({'POST'...})

which is obviously ugly, but I think the proper solution is probably to have your ALLOWED_METHODS = frozenset({... deduce to a final frozenset of a union of Literals rather than using the frozenset itself as a union.

[AlexWaygood]

Closing as a duplicate of #3229