Python 使用嵌套生成器进行操作

Question

我有一些这样的代码：

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)

我想做的是在调用

next(it)

时分解调用链。更具体地说，我想在每个类中的

yield

之后执行一些新代码，而不修改类代码。理想情况下，新代码会打印

yield

在哪个类中执行。这可能吗？

Answer 1

您可以使用包装器

__iter__

方法为每个生成器类创建一个子类，该方法打印在从超类的

yield

方法产生之前要在哪个类

__iter__

中执行。由于

yield

暂停生成器的执行并将控制权返回给调用者，因此包装器应该在产生之前而不是之后打印这样的输出：

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)

演示：https://replit.com/@blhsing/UprightWebbed人工智能