高阶Python｜返回类型提示技巧 (2)

Python提供了一种可选的特性——类型提示，它有助于提高代码的可读性、可推理性和可调试性。通过类型提示，开发者能够清楚地了解变量、函数参数和返回值应具备的数据类型。在开发那些需要高度灵活性的应用程序时，您可能会需要为代码指定多种可能的返回类型，这样做可以使您的程序更加健壮，更能适应多变的运行环境。

在实际开发中，您可能会碰到需要在Python函数中标注多种返回类型的情况。这意味着函数返回的数据类型不是单一的，而是多样的。本文[1]将通过实例向您展示，如何为一个从电子邮件地址中解析出域名的函数定义多种可能的返回类型。同时，您还将学习到如何为那些接受函数作为参数或者作为回调的函数添加类型提示。通过这些示例，您将能够更自如地在函数式编程中使用类型提示。

为工厂函数的返回值添加类型注解

工厂函数是一种特殊的高阶函数，它能够根据给定的参数从头创建一个新的函数。这种工厂函数的参数会影响新创建函数的行为。特别地，在Python中，如果一个函数接收一个可调用对象作为参数，并且返回一个可调用对象，这样的函数被称为装饰器。

延续之前的例子，假设您想要编写一个装饰器来测量代码中其他函数的执行时间。以下展示了如何测量 parse_email() 函数执行所需的时间长度：

>>> import functools
>>> import time
>>> from collections.abc import Callable
>>> from typing import ParamSpec, TypeVar

>>> P = ParamSpec("P")
>>> T = TypeVar("T")
>>> def timeit(function: Callable[P, T]) -> Callable[P, T]:
...     @functools.wraps(function)
...     def wrapper(*args: P.args, **kwargs: P.kwargs):
...         start = time.perf_counter()
...         result = function(*args, **kwargs)
...         end = time.perf_counter()
...         print(f"{function.__name__}() finished in {end - start:.10f}s")
...         return result
...     return wrapper
...

timeit() 装饰器可以接收任意输入输出的可调用对象，并返回一个具有相同输入输出的可调用对象。ParamSpec 注解用来指明 Callable 中的任意输入参数，而 TypeVar 用来表示 Callable 中的任意输出类型。

装饰器内部定义了一个名为 wrapper() 的函数，它利用计时器来测量传入的可调用对象执行所需的时间。wrapper() 函数首先记录当前时间到 start 变量，然后执行被装饰的函数，记录执行完毕后的时间到 end 变量，并计算出执行持续的时间。计算完毕后，它会打印出这个时间，然后返回被装饰函数的执行结果。

定义好 timeit() 装饰器后，您可以使用 @ 符号来简洁地将任何函数应用这个装饰器，而不需要手动调用它，就像使用一个工厂函数那样。例如，您可以将 @timeit 应用于 parse_email() 函数，从而创建一个新函数，这个新函数除了执行原有功能外，还会自动计时自己的执行过程。

>>> @timeit
... def parse_email(email_address: str) -> tuple[str, str] | None:
...     if "@" in email_address:
...         username, domain = email_address.split("@")
...         return username, domain
...     return None
...

您以一种声明性的方式为函数增加了新的能力，而无需更改其源代码，这种做法既优雅又简洁，尽管它可能与Python的设计哲学略有不符。一些人可能会认为装饰器的使用使得代码的直观性有所降低。然而，装饰器同时也能够让代码的外表更加简洁，从而提升代码的易读性。

当您执行经过装饰的 parse_email() 函数时，它不仅会返回预期的结果，还会输出一条信息，这条信息说明了原始函数执行完成所花费的时间。

>>> username, domain = parse_email("claudia@realpython.com")
parse_email() finished in 0.0000042690s

>>> username
'claudia'

>>> domain
'realpython.com'

您函数的执行时间非常短，正如上面的消息所指出的那样。在调用了装饰过的函数之后，您将得到的元组赋值并分解到名为 username 和 domain 的变量里。

为生成器函数返回的值添加类型注解

在某些情况下，为了提高效率，特别是处理大型数据集时，您可能更倾向于使用生成器逐个产生数据片段，而不是将所有数据一次性加载到内存中。在Python中，您可以为生成器函数添加类型提示。一种常见的做法是使用collections.abc模块中的Generator类型进行注解。

请注意，就像Callable类型一样，不要将collections.abc.Generator与已经被弃用的typing.Generator类型相混淆。

以前面的例子为基础，设想您现在需要处理一个很长的电子邮件列表。与其将每个解析结果都存储在内存中，并让函数一次性返回所有结果，不如使用生成器逐个产生解析后的用户名和域名。

为此，您可以编写如下的生成器函数，该函数逐个产生所需信息，并使用Generator类型作为返回值的类型提示：

>>> from collections.abc import Generator

>>> def parse_email() -> Generator[tuple[str, str] | str, str, str]:
...     sent = yield "", ""
...     while sent != "":
...         if "@" in sent:
...             username, domain = sent.split("@")
...             sent = yield username, domain
...         else:
...             sent = yield "invalid email"
...     return "Done"
...

parse_email() 这个生成器函数不需要任何参数输入，因为您将直接向生成器对象发送数据。请注意，Generator 类型注解需要三个参数，其中后两个是可选项：

产生类型：第一个参数定义了生成器将产生什么类型的数据。这里指的是一个元组，包含两个字符串，分别代表从电子邮件地址解析出的用户名和域名。如果电子邮件地址无效，生成器也可能产生一个表示错误的字符串。 发送类型：第二个参数说明了您将向生成器发送什么类型的数据。这同样是一个字符串，因为您将向生成器发送电子邮件地址。 返回类型：第三个参数代表生成器完成所有值的产生后将返回什么。在这个例子中，函数返回的是字符串 "Done"。 以下是您使用这个生成器函数的方法：

>>> generator = parse_email()
>>> next(generator)
('', '')
>>> generator.send("claudia@realpython.com")
('claudia', 'realpython.com')
>>> generator.send("realpython")
'invalid email'
>>> try:
...     generator.send("")
... except StopIteration as ex:
...     print(ex.value)
...
Done

您首先通过调用 parse_email() 生成器函数来启动，这个函数会返回一个新的生成器实例。接着，通过调用内置的 next() 方法，您可以将生成器推进到第一个 yield 语句。从这时起，您就可以开始向生成器发送电子邮件地址，以便进行解析。当您发送一个空字符串时，生成器会停止工作。

由于生成器本质上也是一种迭代器，即生成器迭代器，您也可以选择使用 collections.abc.Iterator 类型作为类型提示，来表达类似的意思。但请注意，如果您的生成器除了产生值之外还有其他操作，比如发送值或返回值，那么使用 collections.abc.Iterator 作为类型提示可能就不够用了，因为它不支持指定发送和返回类型。

from collections.abc import Iterator

def parse_emails(emails: list[str]) -> Iterator[tuple[str, str]]:
    for email in emails:
        if "@" in email:
            username, domain = email.split("@")
            yield username, domain

这个版本的 parse_email() 函数接收一个由字符串组成的列表，并返回一个生成器对象。这个生成器通过 for 循环以延迟执行（惰性）的方式逐个处理列表中的元素。尽管生成器在功能上比迭代器更具体，但迭代器的概念更广泛，也更易于理解，因此使用迭代器作为类型注解是一个合适的选择。

在某些情况下，Python 开发者可能会选择使用更宽松、更通用的 collections.abc.Iterable 类型来为生成器添加类型注解，这样做的好处是它不会暴露生成器的具体实现细节。

from collections.abc import Iterable

def parse_emails(emails: Iterable[str]) -> Iterable[tuple[str, str]]:
    for email in emails:
        if "@" in email:
            username, domain = email.split("@")
            yield username, domain

在这个例子中，您用 Iterable 类型来标注函数的输入参数和返回值，这样做让函数变得更加灵活。函数现在可以接受任何类型的可迭代对象，不仅限于之前的列表形式。

反过来，调用函数的代码不需要关心返回的是生成器还是一系列元素，只要它们能够迭代处理即可。这为函数的实现提供了极大的灵活性，因为您可以在不破坏与调用者通过类型提示建立的约定的情况下，将实现从立即加载的列表更改为按需产生元素的生成器。当您预见到返回的数据量可能会很大，需要使用生成器来处理时，这种灵活性就显得尤为重要。