Question

问题

我有一个迭代器spam，如果它生成很少的项目我想要应用函数foo，否则bar。换句话说，我希望将以下代码转换为生成器：

if len(spam) <= max_size_for_foo:
    foo(spam)
else:
    bar(spam)

max_size_for_foo相对较小，因此如果创建了这个长度的列表或其他可迭代没有问题，但如果spam很长，则不能将其转换为列表（否则，内存）问题随之而来。）

不满意解决方案

到目前为止，我能想出的最佳解决方案如下：

first_items = []
try:
    for i in range(max_size_for_foo+1):
        first_items.append(next(my_generator))
except StopIteration:
    foo(first_items)
else:
    my_generator = chain(first_items, my_generator)
    bar(my_generator)

然而，提取临时列表并链接回生成器对我来说感觉相当肮脏和不优雅。

问题

有没有更优雅或Pythonesque方式来做到这一点？

Answer 1

最简单的方法可能是在函数中定义生成器，以便可以重用它：

def spam_func():
    return (i for i in [1, 2, 3])

spam_length = sum(1 for _ in spam_func())
if spam_length <= max_size_for_foo:
    foo(spam_func())
else:
    bar(spam_func())

Answer 2

没有发电机长度这样的东西，因为发电机很容易无限：

def square():
    a = 0
    while True:
        yield a**2
        a += 1

一个解决方案是sum(1 for _ in gen)。另一个可能是用这样的类包装你的生成器：

class wrapper:
    def __init__(self, items):
        self.items = items

    def __len__(self):
        return len(self.items)

    def generate(self):
        yeild from self.items

根据发生器的长度将不同的功能应用于发生器

问题

不满意解决方案

问题

2 个答案: