Python函数指定整数值的类型提示及mypy检查问题

问题背景 #

我正在学习Python类型提示，遇到一个奇怪的情况：我需要定义一个仅接受0、1、2三个整数作为参数的函数，尝试了两种写法都报错：

1. 直接用字面量作为类型提示：

def func(number: 0 | 1 | 2):
    print(number)

input_number = 0
func(input_number)

执行mypy test.py时出现错误：

test.py:1: error: Invalid type: try using Literal[0] instead?  [valid-type]
test.py:1: error: Invalid type: try using Literal[1] instead?  [valid-type]
test.py:1: error: Invalid type: try using Literal[2] instead?  [valid-type]

2. 改用Literal后：

from typing import Literal

def func(number: Literal[0, 1, 2]):
    print(number)

input_number = 0
func(input_number)

又出现新错误：

test.py:3: error: Argument 1 to "func" has incompatible type "int"; expected "Literal[0, 1, 2]"  [arg-type]

误以为Literal只适用于字符串，改成input_number = '0'后依然报错，只是类型不匹配的提示变成了字符串。

错误原因与解决方案 #

第一个错误的原因

Python的类型提示要求必须使用合法的类型对象，直接用0 | 1 | 2这种字面量组合不是有效的类型，因此mypy提示你改用Literal——这是正确的方向，Literal就是用来标注仅接受特定字面量值的类型。

第二个错误的原因

mypy会自动将input_number = 0推断为int类型，而不是Literal[0]。当你把一个普通int类型的变量传给要求Literal[0,1,2]的函数时，mypy会认为类型不兼容，因为普通int可能包含0、1、2之外的值。

解决方法

有三种常见的解决方式：

1. 给变量添加明确的Literal类型注解
   告诉mypy这个变量的取值范围是指定的字面量：

from typing import Literal

def func(number: Literal[0, 1, 2]):
    print(number)

# 明确标注变量类型为Literal[0]
input_number: Literal[0] = 0
func(input_number)

# 或者标注为允许的所有值
input_number: Literal[0, 1, 2] = 0
func(input_number)

2. 直接传入字面量参数
   如果你直接传入0、1、2这些字面量，mypy会自动识别它们为Literal类型，不会报错：

from typing import Literal

def func(number: Literal[0, 1, 2]):
    print(number)

func(0)  # 正常通过mypy检查
func(1)  # 正常通过mypy检查

3. 使用类型断言（cast）
   如果你确定变量的取值符合要求，可以用cast强制告诉mypy变量的类型：

from typing import Literal, cast

def func(number: Literal[0, 1, 2]):
    print(number)

input_number = 0
func(cast(Literal[0], input_number))

补充说明 #

• Literal并非仅适用于字符串类型，它支持int、str、bool、None等多种字面量类型。
• mypy是静态类型检查工具，不会运行代码，因此它只能通过类型注解来判断值的合法性，无法直接推断未标注类型的变量的具体取值范围。

内容的提问来源于stack exchange，提问作者Abdulrahman Sheikho