浅谈 Python 生成器
Apr 22, 2020
4 minute read

生成器是本人最喜欢的语法糖之一。通过简单 yield 表达式即可实现各种奇妙用法。它解决用 Python 语言实现生产者-消费者模型的一些痛点。在没有生成器语法前,差不多有四种方式来处理生产者-消费者模型问题。

四种实现生产者-消费者模型的方法

  1. 通过传递回调函数,通过回调函数消费每次生产的变量(或者通过回调生产每次需要消费的变量)。

    def countdown(number: int, handler):
        while number > 0:
            number -= 1
            handler(number)
    
    countdown(3, print)
    # 2
    # 1
    # 0

    通过以上简单的例子,似乎这方法没有啥痛点的。 如果是实现语法分析器的话。在把单词流转换为语法树的过程中,回调函数就需要通过全局变量来维护分析器的状态。这样就会导致全局变量满天飞,不仅难以实现正确,而且代码可阅读性也不高。

    这个例子举的不是很恰当,毕竟用生成器实现起来也不简单(难点完全不在这里)。为了避免以上提到的问题,还可以通过闭包或者面向对象编程来解决。

    若通过以下的三个方法中的任意一个去实现的话,则可以就地去做处理,避免了全局变量漫天飞的情况。

  2. 直接等生产结束,返回一个超长的列表,再进行消费。

    def countdown(number: int):
        rv = []
        while number > 0:
            number -= 1
            rv.append(number)
        return rv
    
    countdown(3)
    # [2, 1, 0]

    这例子所存在的问题很容易看出来

    1. 一次性产生完会耗费大量的内存,甚至会导致 OOM;
    2. 无法实现无限制地生产;
    3. 在只需要前几个变量的情况下,一次性生产所有未免太浪费了。不仅会浪费时间,还会浪费内存(可能会造成不必要的内存毛刺)。
  3. 按照 Iterator 迭代器接口的定义类,以类变量来维护自身状态,定义魔法方法 __next__ 来生产变量。 (不一定要按照迭代器接口实现,实现个普通类及普通方法也可以的。但实现迭代器接口的话,还可以使用 for-in 语法糖哟)。

    class countdown:
        def __init__(self, number: int):
            self.number = number
    
        def __next__(self):
            if self.number > 0:
                self.number -= 1
                return self.number
            raise StopIteration
    
        def __iter__(self):
            return countdown(self.number)
    
    for i in countdown(3):
        print(i)
    
    # 2
    # 1
    # 0

    缺点只能说是,实现起来太麻烦了。

  4. 通过使用线程 + 队列,并发地生产及消费。

    import threading
    
    from queue import Queue
    
    def countdown(number: int, q_out: Queue):
        while number > 0:
            number -= 1
            q_out.put(number)
    
        q_out.put(None)
    
    def consumer(q_in: Queue):
        while True:
            v = q_in.get()
            if v is None:
                break
            print(v)
    
    q = Queue()
    thr_consumer = threading.Thread(target=consumer, args=(q, ))
    thr_consumer.start()
    threading.Thread(target=countdown, args=(3, q)).start()
    thr_consumer.join()
    # 2
    # 1
    # 0

    无论是哪种语言,大家首先想到的一定就是这个方法吧。这个方法的缺点也很明显,用多线程来实现生产者-消费者模型,会多出了线程切换的开销。

PEP 255 通过之后,Python 就有了第五种方法 – 生成器。

  • Generator – 生成器函数(类)
    • 只要在函数中使用了 yield 表达式,就可以被称为生成器函数;
    • 只要实现了生成器接口的类,其实例就可以当作生成器来使用。
  • Generator Iterator – 生成器

普遍都直接用 Generator 称呼,但为了避免误会还是要尽量用全称。

yield 关键字

只要在函数中使用 yield 表达式,即可实现一个 Generator 生成器函数。相比较之前的四个方法,是不是显得非常简单且直观。

def countdown(number):
    while number > 0:
        number -= 1
        yield number

countdown(3) 调用后生成 Generator Iterator 生成器,通过 for-in 即可从 2 数到 0。

for number in countdown(3):
    print(number)

# 2
# 1
# 0

生成器的四个状态

Generator Iterator States

  • NotStart 未启动
  • Running 运行
  • Paused 暂停
  • Stopped 完全停止

以上面的 countdown 生成器函数为例,执行 gen = countdown(3) 生成生成器。

当第一次执行 next(gen) 时,生成器从状态 NotStart 转变为状态 Running,即函数开始执行。直到遇到 yield number,再转变为状态 Paused。被 yield 的值 number=2 就是 next(gen) 表达式的值。

Generator Iterator First Next

第二次,第三次执行 next(gen) 时,生成器从状态 Paused 转变为状态 Running,即从中断的地方开始执行,经过 while loops 条件的判断 number > 0True,继续执行 while loops 的代码块。知道遇到 yield number,再转变为状态 Paused。两次 next(gen) 的值分别为 10

Generator Iterator Next

最后一次执行 next(gen) 时,生成器从状态 Paused 转变为状态 Running。在 while loops 条件的判断 number > 0False, 跳出 while loops,没有新的代码可执行,函数就执行结束了。当函数执行结束时,next(gen) 表达式执行没结果,抛出 StopIteration 异常。

Generator Iterator Last Next

搞明白了怎么定义一个生成器函数,及其是怎么运行的。但只把生成器当迭代器用也太……没用了。不要着急,后面会讲到怎么用生成器来整活。

实现个具有生成器接口的类

还可通过继承 collections.abc.Generator 元类(可以不继承)来实现具有生成器接口的类。只要实现了生成器接口的类,其实例就可以当作生成器来使用。

PEP-342 以后,才有生成器接口这个标准。 这个 PEP 扩展了生成器的语法。 提供了让调用者能干预到生成器内部的方法 sendthrow,这样就可以利用生成器做协程。

  • 必需实现的方法

    Generator 接口为 Iterator 接口的子类

    • send(value)
    • throw(type, value=None, traceback=None)

    Iterator 接口

    • __next__() 等同于 send(None),可以通过内建函数 next 去调用。
    • __iter__() 获取可迭代对象,即生成器本身

      def __next__(self):
          return self.send(None)
      
      def __iter__(self):
          return self

    实现了 Iterator 接口的方法,才能用 for-loops 语法。直接复制使用以上这两个方法的代码即可。

  • 可选实现的方法

    • close() 可以直接复制下方的 close 的代码

send

send(value) 用来使生成器继续运行,而 valueyield 表达式的值。

def receiver():
    while True:
        value = yield
        print("receive", value)

recv = receiver()
# equal to `next(recv)` and `recv.__next__()`. Evaluate the first next call.
recv.send(None)
recv.send(1)
# receive 1
recv.send("abc")
# receive abc

throw

throw(type, value=None, traceback=None) 也是用来使生成器继续运行,在生成器停止的位置抛出传递的异常。

def catcher():
    while True:
        try:
            yield
        except Exception:
            print("catch", sys.exc_info())
c =  catcher()
c.send(None)
c.throw(TimeoutError)
c.throw(TimeoutError("boo"))
c.throw(TimeoutError, TimeoutError("boo"))
c.throw(TimeoutError, "boo")
c.throw(TimeoutError, ("boo", 1))
# catch (<class 'TimeoutError'>, TimeoutError(), <traceback object at 0x107335700>)
# catch (<class 'TimeoutError'>, TimeoutError('boo'), <traceback object at 0x107335700>)
# catch (<class 'TimeoutError'>, TimeoutError('boo'), <traceback object at 0x107267bc0>)
# catch (<class 'TimeoutError'>, TimeoutError('boo'), <traceback object at 0x107267bc0>)
# catch (<class 'TimeoutError'>, TimeoutError('boo', 1), <traceback object at 0x107335700>)

这个 throw 方法参数设计真是很奇妙……

  1. 可以直接传递异常类型
  2. 可以接传递异常
  3. 可以传递异常类型加上类型的传参,用 tuple 来作多个传参。(

这个函数参数设计得如此奇妙的原因可能是,因为是历史包裹吧。

以上使用简单生成器函数例子,简明地介绍了生成器接口的两个方法 sendthrow

close

通过调用 close 来完全停止生成器,其定义可以参考下面的语义等同的 Python 代码(原生的肯定是用 C 实现的)。这函数就简单地利用一下 throw 方法在生成器停止的位 yield 表达式位置处抛出 GeneratorExit 异常,使得生成器完全停止。

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

虽说这一节讲的是用类来定义生成器接口,但举的例子都是生成器函数……为什么不直接用类定义来做介绍呢?因为单纯实现个具有生成器接口的类没有什么必要,不直观,而且不好正确地实现。 还不如直接定义生成器函数。所以以生成器函数的方式介绍 sendclose 方法,更直观,更容易理解。但在做元编程,或者实现 c 扩展就非常有用了。这种情况下,只能通过实现具有生成器接口来实现生成器了。

小技巧与细节

Resource Management

举个例子,用生成器函数来实现一个用来过滤重复项的 Pipeline。这就需要一个超大集合用来判断是否重复,可能是通过读写磁盘,或者外部服务等方式实现的,难免会需要在使用结束后做一些清理工作。

def filter_duplicated(sequence):
    big_set = ...
    try:
        for value in sequence:
            if value in big_set:
                continue

            big_set.add(value)
            yield value
    finally:
        big_set.close()

cursor = con.execute()
original_seq = iter(cursor.fetchone, None)
filtered_seq = filter_duplicated(map(lambda row: row[0], original_seq))
for url in filtered_seq:
    ...

若是要在生成器完全停止前做清理,只要在 yield 表达式外面套一层 try-finally。在 finally 代码块做清理工作。当主动调用 close 方法时,会在之前暂停的 yield 表达式位置处抛出 GeneratorExit 异常。因在 try 代码块中抛出,所以一定会执行 finally 代码块。

同样利用生成器做资源回收的 contextlib.contextmanager,简单地通过装饰生成器函数,即可使用 with 语法对资源进行管理。如果对这一部分是如何做到抱有疑问,且感兴趣的话可以自行查看源码哟

from contextlib import contextmanager

@contextmanager
def managed_resource(*args, **kwds):
    # Code to acquire resource, e.g.:
    resource = acquire_resource(*args, **kwds)
    try:
        yield resource
    finally:
        # Code to release resource, e.g.:
        release_resource(resource)

>>> with managed_resource(timeout=3600) as resource:
...     # Resource is released at the end of this block,
...     # even if code in the block raises an exception

Return

当生成器函数结束时,会抛出 StopIteration。若是函数有返回值的话,相当于抛出个带参数的 StopIteration 异常。

def gen_with_return():
    yield
    return 1  # equal to `raise StopIteration(1)`

gen = gen_with_return()
next(gen)
try:
    next(gen)
except StopIteration as exc:
    assert exc.value == 1
else:
    raise AssertError

yield from 表达式

超方便,不需要写重复的代码。看 PEP 380 里的等价代码,就知道要正确的处理子生成器是一件麻烦而又重复的事。接下来举个超简单的例子

def chain(*gens):
    for gen in gens:
        yield from gen
def chain(*gens):
    for gen in gens:
        for v in gen:
            yield v

第二个只能处理只需要调用 next 方法的情况(需要用到 sendthrowclose 的情况根本就没处理)。 还比用 yield from 表达式的例子多了一行……

yield from 不一定只接受 Generator, Iterable 对象也可以。这种情况下,如果调用 sendthrow 可能报错哦~

前面提到在生成器中 return 1 相当于抛出一个带参数的 StopIteration(1)。 如果该生成器被 yield from 的话,则 yield from 表达式的值为 1 哟。

def bottom():
    return (yield)

def middle():
    return (yield from bottom())

def top():
    return (yield from middle())

g = top()
g.send(None)
try:
    g.send(13)
except StopIteration as exc:
    assert exc.value == 13
else:
    raise AssertionError

Auto Close

生成器在其引用计数归零或者被 gc 时,会主动调用 close 方法。但定义具有生成器接口类时,其实例化后的生成器却不会… 因为这一部分逻辑只在原生生成器的 C 代码中定义了,只有通过定义生成器函数的方式,使用时 close 方法才会被主动调用。 解决方法也很简单,通过在类定义中加一行 __del__ = close 即可解决。

总结

生成器函数这个语法糖,通过简单的 yield 语句,就能实现各种各样的功能。不经提高了代码的可阅读性,还让我们少写了很多代码。还能通过 try-finally 去做资源自动回收。而且非常适合做 Pipeline 模式,对数据进行流式处理。

扩展阅读

David Beazley 大佬做了很多关于生成器的分享,推荐先看这一篇 Generator Tricks for Systems Programmers

安利一下 @yihong 在评论推荐了 How the heck does async/await work in Python 3.5? 博文。

参考




comments powered by Disqus