Python itertools—takewhile(): multiple predicates

Question

Suppose a generator yields the below tuples one by one (from left to right)

(1, 2, 3), (2, 5, 6), (3, 7, 10), (4, 5, 11), (3, 5, 15), (4, 5, 9), (4, 6, 12)
...

and suppose I'd like to iterate as long as the predicate is true. Let that predicate be sum(yielded_value) < 20. Then the iterator will stop by (3, 5, 15). I can do it with, say:

list(itertools.takewhile(lambda x: sum(x) < 20, some_generator()))

Question, how do I write a similar expression with two predicates? Suppose I want:

list(itertools.takewhile(lambda x: sum(x) < 20 and first_value_of_tuple > 3, some_generator()))

(which, in this case, stop by (4, 6, 12).)

score 7 · Accepted Answer · answered Dec 09 '15 at 19:34

7

You can access to elements of each tuple with index.

list(itertools.takewhile(lambda x: sum(x) < 20 and x[0] > 3, some_generator()))

answered Dec 09 '15 at 19:34

Ahsanul Haque

10,676
4
41
57

Thanks. Is there a way to point at `next yield` (or `previously yielded value`? – blackened Dec 09 '15 at 19:41
1

AFAIK, you can't look back, unless you store the value, but you can look ahead with `next` for more see this answer: http://stackoverflow.com/questions/2425270/how-to-look-ahead-one-element-in-a-python-generator – Ahsanul Haque Dec 09 '15 at 19:47

score 2 · Answer 2 · answered Dec 09 '15 at 20:37

Since everything in itertools is lazily iterated, and you are using and for two predicates, you can simply use two takewhile iterators. Sometimes I find this more readable than putting both predicates in a single predicate function or lambda:

lessthan20 = itertools.takewhile(lambda x: sum(x) < 20, some_generator())
greaterthan3 = itertools.takewhile(lambda x: x[0] > 3, lessthan20)
list(greaterthan3)

It also makes it so that you don't have a single huge one liner if you need to add even more predicates in the future.

Aske Doerge · Answer 3 · 2015-12-09T19:43:47.077

If you have additional predicates and need to access all the elements of your tuples you can also unpack them in your lambda function:

list(itertools.takewhile(lambda (x, y, z): x+y+z < 20 and x > 3 and y < 7 and z > 1, some_generator()))

This also asserts that all your tuples have length 3. If you get a tuple with 4 values, it fails hard, as opposed to continuing silently. Obviously only useful in some contexts.

Python itertools—takewhile(): multiple predicates

3 Answers3