如何计算被其他代码消费的生成器中的项数

Question

我正在创建一个生成器，这个生成器会被另一个函数使用，但我还是想知道生成了多少个项目：

lines = (line.rstrip('\n') for line in sys.stdin)
process(lines)
print("Processed {} lines.".format( ? ))

我想到的最好办法是用一个类来包装这个生成器，这样可以记录数量，或者也可以把它反过来，用send()方法把东西传进去。有没有一种优雅又高效的方法，可以在你不是使用这个生成器的情况下，查看它在Python 2中生成了多少个项目呢？

编辑：这是我最后得到的结果：

class Count(Iterable):
    """Wrap an iterable (typically a generator) and provide a ``count``
    field counting the number of items.

    Accessing the ``count`` field before iteration is finished will
    invalidate the count.
    """
    def __init__(self, iterable):
        self._iterable = iterable
        self._counter = itertools.count()

    def __iter__(self):
        return itertools.imap(operator.itemgetter(0), itertools.izip(self._iterable, self._counter))

    @property
    def count(self):
        self._counter = itertools.repeat(self._counter.next())
        return self._counter.next()

Answer 1

通常，我会把生成器转换成一个列表，然后计算它的长度。如果你觉得这样做会占用太多内存，那么你自己提到的包装类可能是个不错的选择。不过，这样做也不是太复杂：

class CountingIterator(object):
    def __init__(self, it):
        self.it = it
        self.count = 0
    def __iter__(self):
        return self
    def next(self):
        nxt = next(self.it)
        self.count += 1
        return nxt
    __next__ = next

（最后一行是为了和Python 3.x保持兼容。）

Answer 2

如果你不在乎你正在使用这个生成器，你可以直接这样做：

sum(1 for x in gen)

Answer 3

这里有另一种方法，使用了 itertools.count() 的例子：

import itertools

def generator():
    for i in range(10):
       yield i

def process(l):
    for i in l:
        if i == 5:
            break

def counter_value(counter):
    import re
    return int(re.search('\d+', repr(counter)).group(0))

counter = itertools.count()
process(i for i, v in itertools.izip(generator(), counter))

print "Element consumed by process is : %d " % counter_value(counter)
# output: Element consumed by process is : 6

希望这对你有帮助。