python yield关键字全面解析

你是否曾因处理的数据集过大而内存溢出？你是否曾因为处理各种复杂的函数状态而烦恼？It does help!

本文聚焦yield generator, 帮助你解锁python进阶技法，写出更优雅的程序！

先导概念

为了更好的理解本篇推文的内容，读者必须先深刻理解以下三个概念：List comprehension (列表生成式)，Generator (生成器)，Iterator (迭代器)。

• List comprehension

  List Comprehensions (PEP202)，是python内置的用来生成list的一种快捷高效的方式

[x * x for x in range(1, 11)]
[x * x for x in range(1, 11) if x % 2 == 0]

• Generator

  Generator(PEP255, PEP289)，是列表生成式的一种优化方案，列表生成式的list会直接放在内存中，因此其大小必然会受到内存的限制；而生成器就是为了解决这种资源耗费的情况，能够做到先定义，边循环边计算。

# 注意区分生成器和列表生成式的定义方式
# 生成器是用（）、列表生成式是用[]

>>> g = (x * x for x in range(10))
>>> g
<generator object <genexpr> at 0x0000020A6184D0A0>

# 如果要一个一个打印出来，可以通过next()或则__next__()获得generator的下一个返回值
>>> next(g)
0
>>> next(g)
1
>>> next(g)
4
>>> next(g)
9
>>> g.__next__()
16

# 在实际编程中，我们一般这样使用
for n in g:
    print(n)


## 性能问题
# 前者会现在内存中开辟一片空间建一个list，然后再计算sum;
# 后者通过使用生成器表达式来节省内存
>>> sum([x*x for x in range(10)])
>>> sum(x*x for x in range(10))

• Iterator

  可以直接作用于for循环的对象统称为可迭代对象(Iterable)，包括集合数据类型(如list、tuple、dict、set、str)和 generator (生成器、带yield的generator function)。但是集合数据类型和generator有一个很大的区别：generator可以使用next()不断调用，直至StopIteration。在python中，可以被next()函数调用并不断返回下一个值的对象称为迭代器：Iterator (PEP234)，generator是其中一种功能强大的Iterator.

  PEP255: a Python generator is a kind of Python iterator, but of an especially powerful kind.

>>> from collections import Iterator
>>> isinstance((x for x in range(10)), Iterator)
True

>>> isinstance([], Iterator)
False
>>> isinstance('abc', Iterator)
False
>>> isinstance(iter([]), Iterator)
True
>>> isinstance(iter('abc'), Iterator)
True

 

那么带yield的generator function是什么呢？yield在python中存在哪些作用呢？这就是我们今天推文的主要内容了。

generator function

yield关键字最基础的应用当然是生成器函数了（generator function），我们可以通过函数next()或for循环获取生成器的内容。

def generate_num():
    for i in range(3):
        yield i

next(gen)
Out[1]: 5
next(gen)
Out[2]: 6
gen = generate_num()
next(gen)
Out[3]: 0
next(gen)
Out[4]: 1
next(gen)
Out[5]: 2
next(gen)
Traceback (most recent call last):
  File "/Users/jeffery/miniconda3/envs/contentshare/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3251, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-20-6e72e47198db>", line 1, in <module>
    next(gen)
StopIteration
gen
Out[6]: <generator object generate_num at 0x11ea91bd0>

其精髓在于：它提供了一种可以向调用者返回中间结果的方式，但同时保持函数的local状态，以便函数可以从中断的地方再次恢复从而继续执行 。我们可以再来看一个例子，利用yield关键字实现一个计算斐波那契数列的函数：

def fib():
    a, b = 0, 1
    while 1:
       yield b
       a, b = b, a+b

当调用fib()时，a和b分别被赋值为0和1，然后想调用者返回b(b=1)的值。当调用者再次恢复调用fib()函数时，此时代码记住了上一次的状态（即上一次代码执行到了yield b）并从a, b = b, a=b继续执行，然后进入下一次循环。第二次循环向调用者返回更新后bb=1的值，fib()函数再次中断，等待下一次调用恢复。从调用者的视角来看，fib()就是一个Iterator，但是性能却提升了。因为恢复一个生成器调用会比函数调用更加节省资源。

contextmanager

在我们之前的推文《python面向对象编程》中，我们简单介绍过使用__enter__()和__exit__()创建一个具有会话管理/上下文管理器的自定义类。而结合yield我们可以方便地为普通函数注册一个上下文管理器。

from contextlib import contextmanager
from typing import TextIO, Optional


@contextmanager
def open_file(file_name):
    f: Optional[TextIO] = None
    try:
        f = open(file_name, 'r')
        yield f
    finally:
        if f is not None:
            f.close()


with open_file(__file__) as fd:
    print(fd.readline())

## output:
## from contextlib import contextmanager

如下面的例子中，我们通过contextmanager装饰器(如果你对装饰器感兴趣，也许你可以参考我之前在csdn上的推文：jeffery0207 python装饰器详细剖析 将open_file()函数变为一个具有上下文管理器功能的生成器函数。

yield from

yield from对应着PEP380新提出的一个概念，叫委派生成器。其基本用法是：yield from <expr>， <expr>表达式值应该为一个iterable对象。

  我们先来看一个嵌套序列展开示例。比如，我们定义一个嵌套的list: guys = ['lily', 'alen', ['john', 'tom', 'jeffery'], ['chris', 'amy']]，我们需要将其展开为单个元素的形式：

from typing import Iterable


def flatten(items, ignore_types=(str, bytes)):
    for x in items:
        if isinstance(x, Iterable) and not isinstance(x, ignore_types):
            yield from flatten(x)
        else:
            yield x


guys = ['lily', 'alen', ['john', 'tom', 'jeffery'], ['chris', 'amy']]
for x in flatten(guys):
    print(x, end='\t')


## output:
## 1    2    3    4    5    6    7    8

  在上面的例子中，flatten(guys)被称为delegating generator， flatten(x)被称为subgenerator。委派生成器的含义就是，让一个生成器 (delegating generator) 将其部分操作委托给另一个生成器 (subgenerator)。这使得包含 yield 的一段代码可以被分解出来，放在另一个生成器中。此外，subgenerator的返回值将提供给delegating generator。

coroutine

协程(线程)，这属于一个独立的概念和技术方向了，我们这里在这里仅做必要的介绍，如果大家感兴趣我们可以专门出一期推送来分享协程。

    首先**什么是协程**？协程是一种用户态的轻量级线程，允许程序执行被挂起，同时在恰当的时机下被恢复。**线程又是什么**？线程是进程的一个实体，是CPU调度和分派的最小单位。那**进程又是什么**？进程是计算机执行任务的实体，是进行资源分配的最小单位。他们之间的关系是：一个进程可以包含多个线程，一个线程可以包含多个协程；但是协程既不是线程也不是进程，协程是一个特殊的函数。如果你对python 多线程、多进程感兴趣，可以阅读我之前在csdn的推文：[Python Threading 多线程编程](https://blog.csdn.net/jeffery0207/article/details/82716640)，[python mutilprocessing多进程编程](https://blog.csdn.net/jeffery0207/article/details/82958520)。

    在python中实现协程，我们可以借助标准库`asyncio`，协程的本质是实现一个时间循环，将多个协程函数或称之为任务放到事件循环中，事件循环则会循环执行这些任务。当然，我们今天的重点在于，如何通过`yield`关键字实现简单的协程。我们来看一个[利用协程实现并发示例](https://python3-cookbook.readthedocs.io/zh_CN/latest/c12/p12_using_generators_as_alternative_to_threads.html)：

我们首先定义两个生成器函数：

# 定义两个生成器函数 `countdown` and `countup`
# 这两个生成器函数之间互不干扰

def countdown(n):
    while n > 0:
        print('T-minus', n)
        yield
        n -= 1
    print('Blastoff!')


def countup(n):
    x = 0
    while x < n:
        print('Counting up', x)
        yield
        x += 1

因为协程的编程模型是事件循环，所有我们需要再实现一个简单的任务调度器，通过任务调度器并发调度多个生成器函数 (任务)。

from collections import deque

class TaskScheduler:
    def __init__(self):
        self._task_queue = deque()  # 构建双向对象

    def new_task(self, task):
        self._task_queue.append(task)

    def run(self):
        while self._task_queue:
            task = self._task_queue.popleft()
            try:
                # Run until the next yield statement
                next(task)
                self._task_queue.append(task)
            except StopIteration:
                # Generator is no longer executing
                pass

# Example use
sched = TaskScheduler()
sched.new_task(countdown(10))
sched.new_task(countdown(5))
sched.new_task(countup(15))
sched.run()


## output:
T-minus 3
T-minus 2
Counting up 0
T-minus 2
T-minus 1
Counting up 1
T-minus 1
Blastoff!
Counting up 2
Blastoff!
Counting up 3

在上面的例子中，我们实际上已经实现了一个“操作系统”的最小核心部分。 生成器函数就是任务，而yield语句是任务挂起的信号。 调度器循环检查任务列表直到没有任务要执行为止。

    **PEP342**进一步提出将`yield`从一个关键字(statement)变为表达式(expression)，并为生成器增加了几个新的方法：`send()`，`throw()`，`close()`并允许`yield`与`try/finally`联用。这段话信息量很大，我们通过一个[回文数字判断示例](https://realpython.com/introduction-to-python-generators/#using-advanced-generator-methods)来看一下：

def is_palindrome(num):
    """
    普通函数，判断数字是否是回文序列，如1221， 3443
    :param num:
    :return:
    """
    if num // 10 == 0:
        return False
    temp = num
    reversed_num = 0

    while temp != 0:
        reversed_num = (reversed_num * 10) + (temp % 10)
        temp = temp // 10

    if num == reversed_num:
        return True
    else:
        return False


def infinite_palindromes():
    num = 0
    while True:
        if is_palindrome(num):
            i = (yield num)
            if i is not None:
                num = i
        num += 1

  我们首先定义了一个判断回文数字的普通函数is_palindrome，函数具体内容就不展开，读者有兴趣可以自己分析一下其中的数学知识。简而言之，当传入数字是回文数字时，is_palindrome函数返回True；反之，返回False。接着我们定义了一个生成器函数infinite_palindromes，该函数包含了PEP342的一个新特性：将yield从一个关键字(statement)变为表达式(expression)。当其变为表达式之后，具有如下特性：

• yield num表达式的值将被赋值给i，在生成器函数内可以对i进行进一步的操作；当用next调用时，yield num表达式的值为None;

PEP342同时新增了三个方法：

• 新增send(value)方法可以唤醒generator，并将value传送进去作为yield表达式的值; 该方法返回 generator的下一个值；

• 新增throw(Exception)方法用法抛出异常 (Exception);

• 新增close()方法用于关闭generator. close函数定义等同如下代码：

def close(self):
    try:
        self.throw(GeneratorExit)
    except (GeneratorExit, StopIteration):
        pass
    else:
        raise RuntimeError("generator ignored GeneratorExit")
    # Other exceptions are not caught

有了PEP342新增的特性，我们可以再来实现一个更加复杂有趣的协程示例：

from collections import deque


class ActorScheduler:
    def __init__(self):
        self._actors = {}
        self._msg_queue = deque()

    def new_actor(self, name, actor):
        self._msg_queue.append((actor, None))
        self._actors[name] = actor

    def send(self, name, msg):
        actor = self._actors.get(name)
        if actor:
            self._msg_queue.append((actor, msg))

    def run(self):
        while self._msg_queue:
            actor, msg = self._msg_queue.popleft()
            try:
                actor.send(msg)
            except StopIteration:
                pass
            finally:
                print('invoked %s' % actor)

# Example use
if __name__ == '__main__':
    def printer():
        while True:
            msg = yield  # 等待msg
            print('Got:', msg)

    def counter(sched: ActorScheduler):
        while True:
            n = yield  # 等待n
            if n == 0:
                break
            # Send to the printer task
            sched.send('printer', n)
            # Send the next count to the counter task (recursive)
            sched.send('counter', n - 1)


    sched = ActorScheduler()
    sched.new_actor('printer', printer())
    sched.new_actor('counter', counter(sched))

    # Send an initial message to the counter to initiate
    sched.send('counter', 2)
    sched.run()

## Output 能体会理清这个执行过程，就代表你协程真正入门啦~
invoked <generator object printer at 0x10dc1ec00>
invoked <generator object counter at 0x10dc1ece0>
invoked <generator object counter at 0x10dc1ece0>
Got: 2
invoked <generator object printer at 0x10dc1ec00>
invoked <generator object counter at 0x10dc1ece0>
Got: 1
invoked <generator object printer at 0x10dc1ec00>
invoked <generator object counter at 0x10dc1ece0>

在上面的例子中，通过sched.send('counter', 2)把消息注册到任务中，sched.run()启动事件循环，直至任务完成。

 

好啦，以上就是这篇推文的全部内容，基于yield关键字生成器对于实现复杂状态维持、程序内存优化有着非常大的优势。