Python -- grouping consecutive elements satisfying a predicate

Question

I would like to group consecutive elements satisfying a predicate. A doctest for such a function would look something like

>>> group([1, 2, 3, 0, 4, 5, 0, 0, 6], lambda x: x != 0)
    [[1, 2, 3], [4, 5], [6]] 
>>> group([1, 2, 3, 0, 4, 5, 0, 0, 6], lambda x: x == 0)
    [[0], [0, 0]]

I have written a prototype where I use itertool's takewhile, but it's ugly because I keep casting between list and iter. I also don't want to stick to reading the indices of a list because it feels inefficient. Can someone point me how to mix and match the itertools together in the right way?

from itertools import takewhile

def group(l, p):
    blocks = []
    while True:
        i = iter(l)
        taken = list(takewhile(p, i))
        l = list(i)
        if len(taken) > 0:
            blocks.append(taken)
        if len(l) == 0:
            return blocks

Thanks!

Answer 1

You could use itertools.groupby:

In [22]: import itertools as IT

In [23]: [list(g) for k, g in IT.groupby(
          [1, 2, 3, 0, 4, 5, 0, 0, 6], lambda x: x != 0) if k]
Out[23]: [[1, 2, 3], [4, 5], [6]]