Search code examples
pythongeneratorpython-itertools

Output the result of an input generator and then the result of a translation of the result of the generator


The following code:

def test(x):
    for i in x:
        yield i
        i = list(i)
        i[1] = "X"
        yield tuple(i)

list(test(it.product(["A", "B"], ["C"])))

Outputs the following list:

[('A', 'C'), ('A', 'X'), ('B', 'C'), ('B', 'X')]

How would I adapt the function such that the input generator results are listed first and then the results of the translation?

So:

[('A', 'C'), ('B', 'C'), ('A', 'X'), ('B', 'X')]

Solution

  • def test(x):
        x = list(x)
        yield from x
        for i in x:
            i = list(i)
            i[1] = "X"
            yield tuple(i)
    

    or

    def test(x):
        x1, x2 = it.tee(x)
        yield from x1
        for i in x2:
            i = list(i)
            i[1] = "X"
            yield tuple(i)
    

    Attempt This Online!