Mypy can't infer the type of items from a list of Literals

user_5973770 0
Viewed 450

I have a variable x and a list of literals (say 0, 1, 2). I want to convert x to one of those literals: if x is in the list, I return it; otherwise I return a fallback value:

from typing import Literal, Set

Foo = Literal[0, 1, 2]
foos: Set[Foo] = {0, 1, 2}
 
def convert_to_foo(x: int) -> Foo:
  if x in foos:
    # x must have type Foo, yet this doesn't type check
    y: Foo = x
    return y
  return 0

Unfortunately, this does not type check. Mypy returns the following message (see gist):

main.py:9: error: Incompatible types in assignment (expression has type "int", variable has type "Union[Literal[0], Literal[1], Literal[2]]")

If I belong to a list of Foos, then I am a Foo, right? I can't find an answer in the doc, can someone point me in the right direction?

mypy typing python
3 Answers

Excellent question. I think that cast might be the only way to go here:

from typing import Literal, Set, cast

Foo = Literal[0, 1, 2]
foos: Set[Foo] = {0, 1, 2}
 
def convert_to_foo(x: int) -> Foo:
    if x in foos:
        y: Foo = cast(Foo, x)
        return y
    return 0
user_44513150

I have tested your code on windows Python3.8 and did not find any problem.

from typing import Literal, Set

Foo = Literal[0, 1, 2]
foos: Set[Foo] = {0, 1, 2}

def convert_to_foo(x: int) -> Foo:
    if x in foos:
        # x must have type Foo, yet this doesn't type check
        y: Foo = x
        print(y)
        return y
    return 0

>>> convert_to_foo(3)
0
>>> convert_to_foo(2)
2
2

what exactly do you want to do? Here it checks if for 3 and returns 0 for 2 returns 2. Is not that what you check for?

user_72412870

I tried this x:Any, default: int

from typing import Literal, Set
from typing import *

Foo = Literal[0, 1, 2]
foos: Set[Foo] = {0, 1, 2}

def convert_to_foo(x:Any, default: int) -> Foo:
  if x in foos:
    # x must have type Foo, yet this doesn't type check
    y: Foo = x
    return y
  return 0

Success: no issues found in 1 source file

user_72412870
Related