Python类型注解，你需要知道的都在这里了

2 年前

杜赛

在 Python 中定义函数非常简单，像这样：

def say(name):
    return f'Hello {name}!'

但是，有时候也会看到这样的代码：

def say_hi(name: str) -> str:
    return f'Hello {name}!'

函数定义似乎变得复杂些了：多出来这些 str 、 -> 都是什么意思？有什么作用？

本文将由浅入深，好好聊聊 Python 3.5 之后的 类型注解 。理解它将非常有益于优化你的代码。

变量注解

Python 是 动态语言 ，其显著特点是在声明变量时，你不需要显式声明它的类型。

比如这个：

age = 20
print('The age is: ', age + 1)
# Output:
# The age is:  21

你看，虽然代码里没有明确指定 age 的类型，但是程序运行时 隐式推断 出它是 int 类型，因此可以顺利执行 age + 1 的动作。

除此之外，已经确定类型的变量，可以随时更改其类型，比如：

age = 20
print(type(age))
# Output: <class 'int'>
age = '20'
print(type(age))
# Output: <class 'str'>

Python 这种动态特性的好处是它非常的自由，大部分时候你不用纠结类型声明、类型转化等麻烦事，可以用很少的代码完成各种骚操作。但是缺点也在这里：如果你代码某些变量的类型有错，编辑器、IDE等工具无法在早期替你纠错，只能在程序运行阶段才能够暴露问题。

比如下面这个例子：

age = 20
# ...
# 这里进行了一大串的其他指令
# 然后你忘记了 age 应该是 int
# 错误地将其赋值为字符串
age = '20'
print('The age is: ', age + 1)
# Output: TypeError: can only concatenate str (not "int") to str

在项目代码逐渐膨胀之后，上面这种看似弱智的情况可能会经常发生。

因此，Python 3.5 之后引入了 类型注解 ，其作用就是让你可以明确的声明变量的类型，使代码不再那么的自由（放飞自我）。

类型注解还在快速发展中，因此尽量用较新的 Python 版本去尝试它。

比如上面的代码，就可以用类型注解改写了：

age: int = 20

看似多此一举，但是编辑器可以凭借此，找出你那些错误的骚操作。

比如笔者用的 VS Code ，安装好类型注解插件 Pylance 后，如果写出下面的代码：

age: int = 20
age = '20'

那么编辑器会用醒目的方式告诉你：孙子，你这里的类型写错了！（见下图）

Pylance 默认关闭了类型检查，你得在设置中手动打开。其他的编辑器/IDE（比如 Pycharm）也都提供了类似的类型检查，放心用吧。

很简单，但却带来了 巨大的好处 ：

编辑器可以替你揪出代码中关于类型的错误，避免了程序运行过程中各种奇奇怪怪的 Bug 。
在你编写代码时，编辑器可以提示你对象的类型，免得你或者团队成员忘记了。（程序员通常记性不好）。

注意，类型注解仅仅是提供给编辑器进行类型检查的机会，也就是起提示的作用，对 Python 程序的运行不会产生任何影响。也就是说，Python 跟以前一样自由，即使你进行了错误的类型赋值，只要不直接引发错误，程序依旧可以运行。

最后，Python 中几种基本的变量类型都得到了支持：

a: int = 3
b: float = 3.14
c: str = 'abc'
d: bool = False

很简单吧。让我们继续。

函数注解

文章开头提到的那个例子，就是简单的函数类型注解：

def say_hi(name: str) -> str:
    return f'Hello {name}!'

你可以很清楚的知道，这个函数应该接收一个字符串参数 name ，并且返回值应该也是字符串。

带默认值的函数像这样书写：

def add(first: int = 10, second: float = 5.5) -> float:
    return first + second

如果函数没有返回值，那么下面两种写法都可以：

def foo():
def bar() -> None:
    pass

自定义的对象也没问题，像下面这样：

class Person:
    def __init__(self, name: str):
        self.name = name
def hello(p: Person) -> str:
    return f'Hello, {p.name}'

如果要避免循环导入或者注解早于对象定义的情况，可以用字符串代替类型：

def hello(p: 'Person') -> str:
    return f'Hello, {p.name}'
class Person:
    def __init__(self, name: str):
        self.name = name

效果是相同的。

相比变量类型注解，函数里的类型注解更加有用，并且可能是你最频繁用到注解的地方了。

容器类型

列表、字典、元组等包含元素的复合类型，用简单的 list，dict，tuple 不能够明确说明内部元素的具体类型。

因此要用到 typing 模块提供的 复合注解 功能：

from typing import List, Dict, Tuple
# 参数1: 元素为 int 的列表
# 参数2: 键为字符串，值为 int 的字典
# 返回值: 包含两个元素的元组
def mix(scores: List[int], ages: Dict[str, int]) -> Tuple[int, int]:
    return (0, 0)

如果你用的是 Python 3.9+ 版本，甚至连 typing 模块都不需要了，内置的容器类型就支持了复合注解：

def mix(scores: list[int], ages: dict[str, int]) -> tuple[int, int]:
    return (0, 0)

在某些情况下，不需要严格区分参数到底是列表还是元组（这种情况还蛮多的）。这时候就可以将它们的特征抽象为更泛化的类型（泛型），比如 Sequence（序列）。

下面是例子：

# Python 3.8 之前的版本
from typing import Sequence as Seq1
def foo(seq: Seq1[str]):
    for item in seq:
        print(item)
# Python 3.9+ 也可以这么写
from collections.abc import Sequence as Seq2
def bar(seq: Seq2[str]):
    for item in seq:
        print(item)

例子中函数的参数不对容器的类型做具体要求，只要它是个序列（比如列表和元组）就可以。

类型别名

有时候对象的类型可能会非常复杂，或者你希望给类型赋予一个有意义的名称，那么你可以这样定义类型的别名：

from typing import Tuple
# 类型别名
Vector2D = Tuple[int, int]
def foo(vector: Vector2D):
    print(vector)
foo(vector=(1, 2))
# Output: (1, 2)

Vector2D 这个名称清晰的表达了这个对象是一个二维的向量。

与类型别名有点类似的，是用 NewType 创建自定义类型：

from typing import NewType
from typing import Tuple
# 创建新类型
Vector3D = NewType('Vector3D', Tuple[int, int, int])
def bar(vector: Vector3D):
    print(vector)

乍一眼看起来与前面的类型别名功能差不多，但不同的是 NewType 创建了原始类型的“子类”：

# 类型检查成功
# 类型别名和原始类型是等价的
foo(vector=(1, 2))
# 类型检查失败
# NewType创建的是原始类型的“子类”
bar(vector=(1, 2, 3))
# 类型检查成功
# 传入参数必须是 Vector3D 的“实例”
v_3d = Vector3D((4, 5, 6))
bar(vector=v_3d)

具体用哪种，得根据情况而定。

更多类型

NoReturn

如果函数没有返回值，那么可以这样写：

from typing import NoReturn
def hello() -> NoReturn:
    raise RuntimeError('oh no')

注意下面这样写是错误的：

def hello() -> NoReturn:
    pass

因为 Python 的函数运行结束时隐式返回 None ，这和真正的无返回值是有区别的。

Optional

猜猜下面的类型注解错在哪里：

def foo(a: int = 0) -> str:
    if a == 1:
        return 'Yeah'

答案：函数既有可能返回 None ，也有可能返回 str 。单凭返回值注解为 str 是不能准确表达此情况的。

这种“可能有也可能没有”的状态被称为 可选值 ，在某些项目中非常常见。比如 web 应用中某个函数接受账号和密码作为参数，如果匹配则返回用户对象，若不匹配则返回 None 。

因此，有专门的可选值类型注解可以处理这种情况：

from typing import Optional
def foo(a: int = 0) -> Optional[str]:
    if a == 1:
        return 'Yeah'

Union

比 Optional 涵盖面更广的是 Union 。

如果函数的返回值是多种类型中的一种时，可以这样写：

from typing import Union
def foo() -> Union[str, int, float]:
    # ....
    # some code here

上面这个函数可以返回字符串、整型、浮点型中的任意一种类型。

可以发现 Optional 实际上是 Union 的特例： Optional[X] 和 Union[X, None] 是等价的。

Callable

我们知道， Python 中的函数和类的区别并不明显。只要实现了对应的接口，类实例也可以是 可调用的 。

如果不关心对象的具体类型，只要求是可调用的，那么可以这样写：

from typing import Callable
class Foo:
    def __call__(self):
        print('I am foo')
def bar():
    print('I am bar')
def hello(a: Callable):
# 类型检查通过
hello(Foo())
# 同样通过
hello(bar)

Literal

即字面量。它在定义简单的枚举值时非常好用，比如：

from typing import Literal
MODE = Literal['r', 'rb', 'w', 'wb']
def open_helper(file: str, mode: MODE) -> str:
open_helper('/some/path', 'r')  # 成功
open_helper('/other/path', 'typo')  # 失败

Protocol

即协议。我们通常说一个对象遵守了某个”协议“，意思是这个对象实现了”协议“中规定的属性或者方法。

比如下面这个例子：

from typing import Protocol
class Proto(Protocol):
    def foo(self):
        print('I am proto')
class A:
    def foo(self):
        print('I am A')
class B:
    def bar(self):
        print('I am B')
def yeah(a: Proto):
# 通过，A 实现了协议中的 foo() 方法
yeah(A())
# 不通过，B 未实现 foo()
yeah(B())

Any

如果你实在不知道某个类型注解应该怎么写时，这里还有个最后的逃生通道：

from typing import Any
def foo() -> Any:
    pass

任何类型都与 Any 兼容。当然如果你把所有的类型都注解为 Any 将毫无意义，因此 Any 应当尽量少使用。

泛型

要理解泛型，首先得知道没有它时所遇到的麻烦。

假设有一个函数，要求它既能够处理字符串，又能够处理数字。那么你可能很自然地想到了 Union 。

于是写出第一版代码：

from typing import Union, List
U = Union[str, int]
def foo(a: U, b: U) -> List[U]:
    return [a, b]

这样写有个很大的弊端，就是参数的类型可以混着用（比如 a: int 且 b:str ），即便你并不想具有这样的特性。看下面这个就明白了：

# 类型检查通过
# 因为 Union[str, int] 可以是其中任意一种类型
# 即便你并不想将 str 和 int 混用
foo('Joe', 19)
foo(19, 21)
foo('Joe', 'David')

泛型就可以解决此问题。来看第二版代码：

from typing import TypeVar, List
# 定义泛型 T
# T 必须是 str 或 int 其中一种
T = TypeVar('T', str, int)
def bar(a: T, b: T) -> List[T]:
    return [a, b]
# 类型检查不通过
# 函数的参数必须为同一个类型"T"
bar('Joe', 19)
bar(19, 21)
bar('Joe', 'David')

可以看出，泛型类似于某种模板（或者占位符），它可以很精确地将对象限定在你真正需要的类型。

让我们再看看下面这个对泛型的应用：

from typing import Dict, TypeVar
# 定义泛型 K 和 V
# K 和 V 的具体类型没有限制
K = TypeVar("K")
V = TypeVar("V")
def get_item(key: K, container: Dict[K, V]) -> V:
    return container[key]
dict_1 = {"age": 10}
dict_2 = {99: "dusai"}
print(get_item("age", dict_1))
# 类型检查通过，输出: 10
print(get_item(99, dict_2))
# 类型检查通过，输出: dusai
print(get_item("name", dict_2))