Subclassing and overriding a generator function in python

Question

I need to override a method of a parent class, which is a generator, and am wondering the correct way to do this. Is there anything wrong with the following, or a more efficient way?

class A:
    def gen(self):
        yield 1
        yield 2

class B(A):
    def gen(self):
        yield 3
        for n in super().gen():
            yield n

Answer 1

For Python 3.3 and up, the best, most general way to do this is:

class A:
    def gen(self):
        yield 1
        yield 2

class B(A):
    def gen(self):
        yield 3
        yield from super().gen()

This uses the new yield from syntax for delegating to a subgenerator. It's better than the other solutions because it's actually handing control to the generator it delegates to; if said generator supports .send and .throw to pass values and exceptions into the generator, then delegation means it actually receives the values; explicitly looping and yielding one by one will receive the values in the gen wrapper, not the generator actually producing the values, and the same problem applies to other solutions like using itertools.chain.

Answer 2

What you have looks fine, but is not the only approach. What's important about a generator function is that it returns an iterable object. Your subclass could thus instead directly create an iterable, for example:

import itertools

class B(A):
    def gen(self):
        return itertools.chain([3], super().gen())

The better approach is going to depend on exactly what you're doing; the above looks needlessly complex, but I wouldn't want to generalize from such a simple example.

Answer 3

To call a method from a subclass you need the keyword super.

New Source Code:

class B(A):
    def gen(self):
        yield 3
        for n in super().gen():
            yield n

This:

b = B()
for i in b.gen():
     print(i)

produces the output:

   3
   1
   2

In the first Iteration your generator stops at '3', for the following iterations it just goes on as the superclass normally would.

This Question provides a really good and lengthy explanation of generators, iterators and the yield- keyword: What does the "yield" keyword do in Python?