CODEWITHSUNDEEP
pythonfunctional-programmingiteratorscopefunction-composition
dividebyzero
dividebyzero

Reputation: 1283

Python: compose list from nested iterators

I have a list of tuples that I need to expand it by adding two elements, each of them comes from a list. So I have [(1, 2, 3)] and two iterators it1 = (i for i in ['a1', 'a2']) and it2 = (i for i in in ['b1', 'b2']). The result should be [(1, 2, 3, 'a1', 'b1'), (1, 2, 3, 'a1', 'b2'), (1, 2, 3, 'a2', 'b1'), (1, 2, 3, 'a2', 'b2')].

If I use iterators as show above it doesn't work. But if I use lists it works. Here is the code:

def get_iters():
    return ((i for i in ['a1', 'a2']), (i for i in ['b1', 'b2']))

def get_lists():
    return ([i for i in ['a1', 'a2']], [i for i in ['b1', 'b2']])

def compose(lst_of_tuples, iter=True):
    iters = get_iters() if iter else get_lists()
    for it in iters:
        lst_of_tuples = [t + (i,) for t in lst_of_tuples for i in it]
    return lst_of_tuples

print compose([(1,2,3)], True)
# WRONG!???? (what happened to the 'a2' part of it?)
# prints: [(1, 2, 3, 'a1', 'b1'), (1, 2, 3, 'a1', 'b2')]

print compose([(1,2,3)], False)
# RIGHT!! prints: [(1, 2, 3, 'a1', 'b1'), (1, 2, 3, 'a1', 'b2'), (1, 2, 3, 'a2', 'b1'), (1, 2, 3, 'a2', 'b2')]

I can't think of a reason why that would be. Can someone explain?

Upvotes: 2

Views: 1775

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1123520

Iterables can only be iterated over once, after which they are exhausted.

When looping over a given iterable a second time in a for loop, no more elements are returned.

Loop over itertools.product() first instead, then over your list of tuples to generate ouput:

from itertools import product

def compose(lst_of_tuples, iter=True):
    iters = get_iters() if iter else get_lists()
    return [t + i for i in product(*get_iters()) for t in lst_of_tuples]

This produces:

>>> print compose([(1,2,3)], True)
[(1, 2, 3, 'a1', 'b1'), (1, 2, 3, 'a1', 'b2'), (1, 2, 3, 'a2', 'b1'), (1, 2, 3, 'a2', 'b2')]
>>> print compose([(1,2,3)], False)
[(1, 2, 3, 'a1', 'b1'), (1, 2, 3, 'a1', 'b2'), (1, 2, 3, 'a2', 'b1'), (1, 2, 3, 'a2', 'b2')]

Upvotes: 5

Related Questions