Python maneuvering with nested generators

Question

I have some codes like this:

class A:
  def __iter__(self):
    for i in range(100):
      yield i

class B:
  def __init__(self, src):
    self.src = src
  def __iter__(self):
    for i in self.src:
      yield i * 2

class C:
  def __init__(self, src1, src2):
    self.src1 = src1
    self.src2 = src2
  def __iter__(self):
    for i, j in zip(self.src1, self.src2):
      yield i + j

a = A()
b = B(a)
c = C(a, b)
it = iter(c)
next(it)

What I want to do is to decompose the calling chain when invoking next(it). More specifically, I want to execute some new code after the yield in each class without modifying the class codes. Ideally the new code prints in which class the yield is executed. Is this a possible thing?

Answer 1

You can create a subclass for each of the generator class with a wrapper __iter__ method that prints in which class yield is to be executed before yielding from the __iter__ method of the superclass. Since yield pauses the execution of the generator and returns control back to the caller, the wrapper should print such an output before, not after, yielding:

def print_yields(cls):
    class wrapper(cls):
        def __iter__(self):
            iterator = super().__iter__()
            for i in iterator:
                print(f'yielding {i} from class {cls.__name__} with iterator {id(iterator):x}')
                yield i
    return wrapper

a = print_yields(A)()
b = print_yields(B)(a)
c = print_yields(C)(a, b)
for i in c:
    print(i)

Demo: https://replit.com/@blhsing/UprightWebbedArtificialintelligence