Can a Python function take a generator and return generators to subsets of its generated output?

generator, python

Solution

Sure, this does what you want:

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` does the bulk of the work here -- the `key=` just tells it what field to group by.

The resulting generator needs to be wrapped in an `imap` just because you've specified your `printer` signature to take an iterator over number, while, by nature, `groupby` returns iterators over the same items it gets as its input -- here, 2-items tuples with a letter followed by a number -- but this is not really all that germane to your question's title.

The answer to that title is that, yep, a Python function can perfectly well do the job you want -- `itertools.groupby` in fact does exactly that. I recommend studying the itertools module carefully, it's a very useful tool (and delivers splendid performance as well).

Problem

Let's say I have a generator function like this: ``` 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) ``` Example output might be: ('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) I would like to break this into three groups: Group A, Group B, and Group C. And I would like a generator for each group. Then I'd pass the generator and the group letter into a subfunction. An example of the subfunction: ``` def printer(group_letter, generator): print "These numbers are in group %s:" % group_letter for num in generator: print "\t%s" % num ``` The desired output would be: ``` 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 ``` How can I do this without changing big_gen() or printer(), and avoid storing the entire group in memory at once? (In real life, the groups are huge)

Original source