Search code examples
pythongeneratoryield

Concatenate while yielding

How can I concatenate two items when yielding from a function in python?

The base case:

import itertools

def test():
    for number in range(0,10):
        yield number

list(test()) # [0...9]

What if I want to yield both the number and its square number**2

import itertools

def test():
    for number in range(0,10):
        yield itertools.chain.from_iterable([
            number,number*2
        ])

list(test()) 

# [0,0,1,1,2,4,3,9,...] pretended
# <itertools.chain at 0x135decfd0> ... what I got

However doing itertools.chain.from_iterable([generator1, generator2]) from the outside gives the expected result.

def first_list():
    for number in range(0,5):
        yield number

def second_list():
    for number in range(5,10):
        yield number

list(itertools.chain.from_iterable([first_list(), second_list()]))
Solution

  • A simple way is:

    def test():
    for number in range(0,10):
        yield number 
        yield number**2