Yield a python reference multiple times in combination with zip

user_5415668 0
Viewed 46

I have a rather big array in numpy and want to yield the values to the next processing function in the pipeline. I have preallocated numpy arrays, as I don't want to allocate memory every time.

In one function I have to yield two values, the first one should be consumed directly afterwards, the second one a few pipeline stages later.

A simplified version of the code looks like the following:

import numpy as np

def numbers():
    ret = np.array([0])
    for i in range(4):
        ret[0] = i
        yield ret

def raw(gen):
    for d in gen:
        yield d

def add_ten(gen):
    for d in gen:
        yield d, d+10

def fun1():
    nums = numbers()
    g1, g2 = zip(*add_ten(nums))
    for n1, n2 in zip(g1, g2):
        print(n1, n2)

def fun2():
    nums = numbers()
    g = raw(nums)
    for n in g:
        print(n)

print("fun1")
fun1()

print("fun2")
fun2()

The code returns the following

fun1
[3] [10]
[3] [11]
[3] [12]
[3] [13]
fun2
[0]
[1]
[2]
[3]

The generator yields always the same reference, I just changed the underlying data. It's clear, if I would collect the data in a list, I'd see always the same value, as only the reference is stored. However, I don't understand why g1 from fun1 does not work, while g from fun2 works as expected.

Edit: I'd appreciate a workaround for this problem.

Edit: In my application I pass these iterators into functions again and join them after a few steps.

Thierry already had a solution, if I would consume the values directly, thank you for that :)

def fun3():
    nums = numbers()
    g1, g2 = zip(*add_ten(nums))
    h1 = raw(g1)
    h2 = raw(g2)

    for n1, n2 in zip(h1, h2):
        print(n1, n2)
yield iterator python
1 Answers

This is because of the unpacking *add_ten(nums) that takes place in g1, g2 = zip(*add_ten(nums)).

The add_tens generator is immediately exhausted in order to create the unpacked list. This list, that will be the argument for zip, will contain:

(<your array>, <a new array containing 10>),
(<your array>, <a new array containing 11>),
(<your array>, <a new array containing 12>),
(<your array>, <a new array containing 13>),

as yielded by add_ten.

Once zipped, you will have:

(<your array>, <your array>, <your array>, <your array>),
(<a new array containing 10>, ... <a new array containing 13>)

When you print the content of <your_array>, you'll get 3 for each of the 4 references, as this is its current content.

In fun2, as you print the content of your array at each step of the iteration, you will get its successive contents, so 0, 1, 2, 3.

user_5500940
Related