python - Mypy 无法从文字列表中推断出项目的类型

Question

我有一个变量 x 和一个文字列表(比如 0、1、2)。我想将 x 转换为这些文字之一:如果 x 在列表中，我返回它；否则我返回一个后备值:

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

不幸的是，这没有类型检查。 Mypy 返回以下消息(参见 gist ):

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

如果我属于 Foo 的列表，那么我就是一个 Foo，对吗？我在文档中找不到答案，有人可以指出正确的方向吗？

Answer 1

好问题。我认为 cast 可能是到达此处的唯一方法:

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