Is there a way to copy an arbitrary generator in Python?

Viewed 73

Among the best-known features of functional programming are lazy evaluation and infinite lists. In Python, one generally implements these features with generators. But one of the precepts of functional programming is immutability, and generators are not immutable. Just the opposite. Every time one calls next() on a generator, it changes its internal state.

A possible work-around would be to copy a generator before calling next() on it. That works for some generators such as count(). (Perhaps count() is not generator?)

from itertools import count

count_gen = count()
count_gen_copy = copy(count_gen)
print(next(count_gen), next(count_gen), next(count_gen))  # => 0 1 2
print(next(count_gen_copy), next(count_gen_copy), next(count_gen_copy))  # => 0 1 2

But if I define my own generator, e.g., my_count(), I can't copy it.

def my_count(n=0):
    while True:
        yield n
        n += 1


my_count_gen = my_count()
my_count_gen_copy = copy(my_count_gen)
print(next(my_count_gen), next(my_count_gen), next(my_count_gen))
print(next(my_count_gen_copy), next(my_count_gen_copy), next(my_count_gen_copy))

I get an error message when I attempt to execute copy(my_count_gen): TypeError: can't pickle generator objects.

Is there a way around this, or is there some other approach?

Perhaps another way to ask this is: what is copy() copying when it copies copy_gen?

Thanks.

P.S. If I use __iter__() rather than copy(), the __iter__() version acts like the original.

my_count_gen = my_count()
my_count_gen_i = my_count_gen.__iter__()
print(next(my_count_gen), next(my_count_gen), next(my_count_gen))  # => 0 1 2
print(next(my_count_gen_i), next(my_count_gen_i), next(my_count_gen_i))  # => 3 4 5
2 Answers

There's no way to copy arbitrary generators in Python. The operation just doesn't make sense. A generator could depend on all sorts of other uncopyable resources, like file handles, database connections, locks, worker processes, etc. If a generator is holding a lock and you copied it, what would happen to the lock? If a generator is in the middle of a database transaction and you copy it, what would happen to the transaction?

The things you thought were copyable generators aren't generators at all. They're instances of other iterator classes. If you want to write your own iterator class, you can:

class MyCount:
    def __init__(self, n=0):
        self._n = n
    def __iter__(self):
        return self
    def __next__(self):
        retval = self._n
        self._n += 1
        return retval

Some iterators you write that way might even be reasonably copyable. Others, copy.copy will do something completely unreasonable and useless.

While copy doesn't make sense on a generator, you can effectively "copy" an iterator so that you can iterate it many times. The easiest way is to use tee from the itertools module.

def my_count(n=0):
    while True:
        yield n
        n += 1

a, b, c = itertools.tee(my_count(), 3)

# now use a, b, c ...

This uses memory to cache the iterator's results and pass them on.

Related