CODEWITHSUNDEEP
pythongenerator
linkyndy
linkyndy

Reputation: 17928

Loop through generator two items at a time

I have defined a generator that yields log entries from Elasticsearch:

def gen():
    ....
    for hit in results:
        yield hit

How can I loop through two elements at the same time? Something in the lines of:

for one, two in gen():
    ...

By two elements I mean this: A, B, B, C, ..., Y, Z (for a generated list of A, B, ..., Y, Z).

Upvotes: 7

Views: 2414

Answers (3)

Chris_Rands
Chris_Rands

Reputation: 41228

This answer assumes you want non-overlapping pairs. You can do this via zip() because the iterator is consumed:

for one, two in zip(gen, gen):
    # do something

Example:

>>> gen = (x for x in range(5))
>>> for one, two in zip(gen, gen): print(one,two)
... 
0 1
2 3

Note, as timgeb commented, you should use itertools.zip_longest if you have an uneven number of elements and you want the last one with a fill value, for example:

>>> gen = (x for x in range(5))
>>> for one, two in zip_longest(gen, gen): print(one, two)
... 
0 1
2 3
4 None

Upvotes: 10

timgeb
timgeb

Reputation: 78780

Answering your updated question, use itertools.tee to construct a second iterator, advance the second iterator once and discard the result, then loop over both iterators in pairs using zip.

>>> from itertools import tee
>>> it = iter('abc')
>>> it1, it2 = tee(it)
>>> 
>>> next(it2, None)
'a'
>>> for first, second in zip(it1, it2):
...     first, second
... 
('a', 'b')
('b', 'c')

Thank you. This is the cleanest, most "pythonic" way to do it? I feel the solution to this simple problem is to complex.

I don't think there is a cleaner solution. In fact, it's the pairwise recipe from the itertools docs:

def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return izip(a, b)

Upvotes: 7

Sijan Bhandari
Sijan Bhandari

Reputation: 3061

I once solved the similar problem with following :

def gen():
    results = [1, 2, 3, 4, 5]
    for i, v in enumerate(results):
        # if i < len()
        if (i + 1) < len(results):
            yield (v, results[i + 1])
        else:
            yield v


output = gen()

for each in output:
    print(each)

Output will be:

(1, 2)
(2, 3)
(3, 4)
(4, 5)
5

Upvotes: 0

Related Questions