python generator returning ndarray

Question

I'm having the following odd behavior with a generator that returns a numpy array. In the real application, I am using scipy.sparse to implement an iterative linear algebra solver: The below is a stripped down version. I have the following:

import numpy as np

def iterable(n):

  nc = np.array([0,0,0])
  yield nc

  while nc[0] < n:
    nc += 1
    yield nc

  raise StopIteration

When i try and populate a list with the output

iter = iterable(4)

print list(iter)

I get

$ python itertest.py 
[array([4, 4, 4]), array([4, 4, 4]), array([4, 4, 4]), array([4, 4, 4]), array([4, 4, 4])]

i.e., the result of the final 'yield' repeated n times. If I change the yield nc to yield nc.tolist() then I get the expected result. If possible I would like to avoid that as I will end up having to re-make the array. What's going on?

Thanks in advance!

Answer 1

You only create one ndarray. When you do nc += 1, you modify it in place. So the later iterations modify the same array already yielded, and then yield it again.

You can see this by printing out the generated arrays one at a time:

>>> it = iterable(3)
>>> x = next(it)
>>> print x
[0 0 0]
>>> y = next(it)
>>> print y
[1 1 1]
>>> print x   # x changed!
[1 1 1]
>>> x is y
True

You can see that the second iteration changes the object that was yielded on the first iteration, because they are the same object.

To fix it, change your nc += 1 line to nc = nc + 1. This will create a new array on each iteration.