Python函数能否接收生成器并返回其生成输出子集的生成器？

Question

假设我有一个生成器函数，长这样：

import random
def big_gen():
  i = 0
  group = 'a'
  while group != 'd':
    i += 1
    yield (group, i)
    if random.random() < 0.20:
      group = chr(ord(group) + 1)

它的输出可能是这样的： ('a', 1), ('a', 2), ('a', 3), ('a', 4), ('a', 5), ('a', 6), ('a', 7), ('a', 8), ('b', 9), ('c', 10), ('c', 11), ('c', 12), ('c', 13)

我想把这些数据分成三组：A组、B组和C组。然后我想为每一组创建一个生成器。接着，我会把这个生成器和组的字母传递给一个子函数。子函数的例子是：

def printer(group_letter, generator):
  print "These numbers are in group %s:" % group_letter
  for num in generator:
    print "\t%s" % num

我想要的输出结果是：

These numbers are in group a:
1
2
3
4
5
6
7
8
These numbers are in group b:
9
These numbers are in group c:
10
11
12
13

我该怎么做才能不改变big_gen()或printer()，并且避免一次性把整个组的数据都存到内存里？（在实际情况中，这些组的数据是非常庞大的）

Answer 1

你这里有个小问题。你希望printer()这个函数能接收每个组的生成器，但实际上你只有一个生成器在输出所有的组。根据我的理解，你有两个选择：

1) 修改big_gen()函数，让它输出生成器：

import random
def big_gen():
  i = 0
  group = 'a'
  while group != 'd':
    def gen():
        i += 1
        yield i
        if random.random() < 0.20:
            group = chr(ord(group) + 1)
    yield group, gen

 from itertools import imap
 imap(lambda a: printer(*a), big_gen())

2) 修改printer()函数，让它能保持状态，并在组变化时进行识别（这样你就可以保留原来的big_gen()函数）：

def printer(generator):
  group = None
  for grp, num in generator:
    if grp != group:
        print "These numbers are in group %s:" % grp
        group = grp
    print "\t%s" % num

Answer 2

当然，这段代码能实现你想要的效果：

import itertools
import operator

def main():
  for let, gen in itertools.groupby(big_gen(), key=operator.itemgetter(0)):
    secgen = itertools.imap(operator.itemgetter(1), gen)
    printer(let, secgen)

groupby 在这里完成了大部分工作——key= 只是告诉它要根据哪个字段来分组。

生成的结果需要用 imap 包裹起来，因为你指定的 printer 函数需要一个数字的迭代器，而 groupby 本身返回的是它输入的相同项目的迭代器——在这里是包含字母和数字的二元组——不过这和你问题的标题关系不大。

关于标题的问题，答案是：没错，Python 函数完全可以做到你想要的——itertools.groupby 实际上就是这样工作的。我建议你仔细研究一下 itertools 模块，它是一个非常有用的工具（而且性能也很出色）。