从生成器中选择特定项目

Question

给定一个生成器对象，我希望能够只迭代具有定义索引的元素。然而，我有过这样的经历，当生成器产生“重”项目（内存方面）时，与按原始顺序对每个元素执行迭代的情况相比，每个项目的加载需要相当长的时间。这是一个最小的例子，在我的真实情况下，我无法直接访问生成器函数（即

mygen

无法修改）。

import itertools
import numpy as np 

def mygen(n):
    
    for k in range(n):
        yield k 
        

mygenerator=mygen(10)

indices_ls=[0,5,8]
mask      =np.zeros(10,dtype=int)
mask[indices_ls]=1

compressed_generator=itertools.compress(mygenerator, mask)
for i in compressed_generator:
    print(i) 

这个最小的代码可以解决问题，因为只打印选定的元素。但是，有没有办法在不使用

compressed_generator

的情况下做到这一点，即直接迭代原始生成器，但只获取指定的索引？回想一下，生成器函数 (

mygen

) 无法修改。

Answer 1

一个解决方案可能是“修补”生成器正在使用的功能（在这种情况下

range()

，但在您的原始代码中您需要修补您的功能）：

from contextlib import contextmanager


def mygen(n):
    for k in range(n):
        yield k

@contextmanager
def patch_range():
    def my_range(x):
        yield 0
        yield 5
        yield 8

    try:
        original_range = __builtins__.range
        __builtins__.range = my_range
        yield
    finally:
        __builtins__.range = original_range

with patch_range():
    for i in mygen(10):
        print(i)

印花：

0
5
8