Python Generator

我是多无聊居然开始写这种东西…

而且严重暴露了我的基本功不扎实的事实 = =! 平时工作中很少会用到稍微复杂一点的 generator, 基本也就是一个 yield 搞定. 而前天刚好翻出来了教授在 Happy Day 上讲的 Python 高级编程的 slide, 于是今天看了一些 yield 的高级用法.

基本的 Generator

其实都知道为毛要用 Generator, 一般的原因是在并不需要一个完整的列表的时候可以节省内存, 另外个原因是 yield 关键词挺好看的, 可以让代码变得更加 cool 一点. 比如

def split_list(l, limit=100):
    length = len(l)
    return [l[limit*i:limit*(i+1)] for i in range(length//limit+1)]

for sub_list in split_list(walkby_list, limit=10):
'''由于do_something太挫了, 一次只能处理很小量的list,
这里是10, 真够挫的...
所以我们把大的列表给拆了.'''
    do_something(sub_list)

# 发现这样好浪费内存.
# 而且你压根不需要完整的split_list的返回值.
# 于是会有

_marker = object()
def split_list_using_generator(l, limit=100):
    for group in (list(g) for g in izip_longest(*[iter(l)]*n,
                                                fillvalue=_marker)):
    if group[-1] is _marker:
        del group[group.index(_marker):]
    yield group

for sub_list in split_list_using_generator(walkby_list, limit=10):
    do_something(sub_list)

或者稍微高级一点的, 我们来遍历一棵树:

def inorder_taversal(node):
    if not node:
        return
    if node.left_child:
        inorder_taversal(node.left_child)
    do_something_with_node(node)
    if node.right_child:
        inorder_taversal(node.right_child)

inorder_taversal(root)
        
# 这样其实也还行, 不过换个方式会更加好点

def inorder_walk_tree(node):
    if not node:
        return
    if node.left_child:
        for i in inorder_walk_tree(node.left_child):
            yield i
    yield node
    if node.right_child:
        for i in inorder_walk_tree(node.right_child):
            yield i

for i in inorder_walk_tree(root):
    do_something_with_node(i)

# 这样是不是cool多了, 而且处理逻辑还跟walk逻辑分开了,
# 想怎么改就怎么改

嗷这样代码确实好看点了… 至少语法高亮方面, 大段的普通文字里多了一个关键字, 代码瞬间变得更加动人了!

但是generator其实有更高级的用法…

进阶的 Generator

使用了 yield 关键字的函数会变成一个 generator. 实际上 yield 关键字实现的是延迟计算, 也就是说代码执行到 yield 这一行的时候, 函数就很开心地自愿地把控制权给交了出来, 顺便把该返回的值给返回了, 因为它期待着下一次调用 next 的时候会把控制权又转交回给它. 嗯不像某些拿走了就不还的群体(瞬间让我想到了那个算个对数还要提到贫下中农张大爷还是什么的, 还要抨击一顿当时的敏感人物以及赞美一下老毛的数学题)… 所以我们在 yield 上有更多的玩法.

可以试试一个 generator 有哪些给我们玩的方法…

>>> def gen():
>>>     while 1:
>>>         yield 'wahaha'
>>> 
>>> f = gen()
>>> r = [name for name in dir(f) if not name.startswith('__')]
>>> print r
['close', 'gi_code', 'gi_frame', 'gi_running', 'next', 'send', 'throw']

嗯请不要问我为什么取 r 的时候不用 generator … gi 开头的几个方法我们先不管, 先来看 next, send, throw 和 close.

  • next

    我们使用 generator 的时候其实就是语法糖帮忙调用了 next 这个方法. 看名字也知道这个就是取下一次 yield 返回的值了. 所以下面的代码其实是等价的:

      def gen():
          count = 0
          while count < 10:
              yield count
              count += 1
    
      # code1
      for i in gen():
          print i
    
      # code2
      f = gen()
      try:
          while 1:
              print f.next()
      except StopIteration:
          pass
    
      # code1和code2是等价的
    

    而如果我们不使用 yield 来返回, 也可以自己写一个类, 实现 __iter____next__ 或者 next 方法, 这种做法有点蛋疼, 但是有时候写一个类会清晰很多. 同样, next(xxx) 就会去调用 xxx__next__ 方法.

  • send

    这是一个很有意思的方法. 就好像 C 程序中经常会出现 if ((c = xxx()) != null) 这样的语句一样, yield 也是可以写在赋值语句的右边的(总觉得这里有个什么名字的, 是叫右值么? 我都忘光了 = =! ). 因为 yield 不仅是返回值的出口, 同时也是 generator 跟外部交流的入口. 例如这样(教授slide里的原文):

      def grep(pattern):
          r = None
          while 1:
              line = yield r
              r = line if pattern in line else None
    
      >>> g = grep('wahaha')
      >>> g.send(None) # generator没启动的时候只能发送None
      >>> g.send('lalala')
    
      >>> g.send('wahaha lala')
      'wahaha lala'
    

    关键就在于 line = yield r 这句了. 当我们使用 send 方法的时候, 会把 send 的参数发送给 generator, 同时把 line 赋成这个值. 然后返回下一个被 yield 出来的值, 或者没有值了触发一个 StopIteration.

    那么如果有多个 yield 怎么办呢? 答案是按顺序一个一个的赋值, 像这样:

      def many_yield():
          x = None
          y = None
          z = None
          while 1:
              r1 = yield x
              r2 = yield y
              r3 = yield z
              print r1, r2, r3
              yield r1 + r2 + r3
    
      In [1]: f = many_yield()
      In [2]: f.send(None)
      In [3]: f.send('1')
      In [4]: f.send('2')
      In [5]: f.send('3')
      1 2 3
      Out[5]: '123'
      In [6]: f.send('4')
      In [7]: f.send('5')
      In [8]: f.send('6')
      5 6 7
      Out[8]: '567'
    

    嗯… 4被那个 yield r1 + r2 + r3 给吞掉了… 因为他没有赋值给任何一个变量…

    另外, 其实 f.next() 就等价于 f.send(None), next 就是不发任何值给 generator, 然后无情地拿回来他的返回值…

  • throw

    Python 里的异常的运用是非常灵活的, 这货其实就是可以在 generator 运行时抛出一个异常, 然后你来决定怎么处理这个异常. 一个简单的场景是提前终止迭代:

      def gen():
          count = 0
          while count < 1000:
              yield count
              count += 1
    
      def run():
          f = gen()
          try:
              for i in f:
                  if i > 500:
                      f.throw(StopIteration)
                  do_something(i)
          except StopIteration:
              pass
    

    好吧这个例子举得真烂… 要终止随便哪里 raise 个都行… 唔我其实还没想到真正使用他的地方…

  • close

    关闭一个 generator. 对一个已经关闭的 generator 调用 next, send 都会引发 StopIteration. 例如贴教授的原文:

      def grep(pattern):
          try:
              while 1:
                  line = yield
                  if pattern in line:
                      print line
          except GeneratorExit:
              print "Goodbye"
    
      >>> g = grep("python")
      >>> g.next()
      >>> g.send("python generator rocks")
      python generator rocks
      >>> g.close()
      Goodbye
    

    其实我们从代码也可以看出来, close实际上就是 g.throw(GeneratorExit).

更加进阶的 Generator

来自己写一些简单的 coroutine, 下一篇再写… 好困了 ¬ ¬