详解Python的迭代器、生成器以及相关的itertools包

对数学家来说，Python这门语言有着很多吸引他们的地方。举几个例子：对于tuple、lists以及sets等容器的支持，使用与传统数学类似的符号标记方式，还有列表推导式这样与数学中集合推导式和集的结构式(set-builder notation)很相似的语法结构。

另外一些很吸引数学爱好者的特性是Python中的iterator(迭代器)、generator(生成器)以及相关的itertools包。这些工具帮助人们能够很轻松的写出处理诸如无穷序列(infinite sequence)、随机过程(stochastic processes)、递推关系(recurrence relations)以及组合结构(combinatorial structures)等数学对象的优雅代码。本文将涵盖我关于迭代器和生成器的一些笔记，并且有一些我在学习过程中积累的相关经验。
Iterators

迭代器(Iterator)是一个可以对集合进行迭代访问的对象。通过这种方式不需要将集合全部载入内存中，也正因如此，这种集合元素几乎可以是无限的。你可以在Python官方文档的“迭代器类型(Iterator Type)”部分找到相关文档。

让我们对定义的描述再准确些，如果一个对象定义了__iter__方法，并且此方法需要返回一个迭代器，那么这个对象就是可迭代的(iterable)。而迭代器是指实现了__iter__以及next(在Python 3中为__next__)两个方法的对象，前者返回一个迭代器对象，而后者返回迭代过程的下一个集合元素。据我所知，迭代器总是在__iter__方法中简单的返回自己(self)，因为它们正是自己的迭代器。

一般来说，你应该避免直接调用__iter__以及next方法。而应该使用for或是列表推导式(list comprehension)，这样的话Python能够自动为你调用这两个方法。如果你需要手动调用它们，请使用Python的内建函数iter以及next，并且把目标迭代器对象或是集合对象当做参数传递给它们。举个例子，如果c是一个可迭代对象，那么你可以使用iter(c)来访问，而不是c.__iter__()，类似的，如果a是一个迭代器对象，那么请使用next(a)而不是a.next()来访问下一个元素。与之相类似的还有len的用法。

说到len，值得注意的是对迭代器而言没必要去纠结length的定义。所以它们通常不会去实现__len__方法。如果你需要计算容器的长度，那么必须得手动计算，或者使用sum。本文末，在itertools模块之后会给出一个例子。

有一些可迭代对象并不是迭代器，而是使用其他对象作为迭代器。举个例子，list对象是一个可迭代对象，但并不是一个迭代器(它实现了__iter__但并未实现next)。通过下面的例子你可以看到list是如何使用迭代器listiterator的。同时值得注意的是list很好地定义了length属性，而listiterator却没有。

>>> a = [1, 2]
>>> type(a)
<type 'list'>
>>> type(iter(a))
<type 'listiterator'>
>>> it = iter(a)
>>> next(it)
1
>>> next(it)
2
>>> next(it)
Traceback (most recent call last):
 File "<stdin>", line 1, in <module>
StopIteration
>>> len(a)
2
>>> len(it)
Traceback (most recent call last):
 File "<stdin>", line 1, in <module>
TypeError: object of type 'listiterator' has no len()

当迭代结束却仍然被继续迭代访问时，Python解释器会抛出StopIteration异常。然而，前述中提到迭代器可以迭代一个无穷集合，所以对于这种迭代器就必须由用户负责确保不会造成无限循环的情况，请看下面的例子：
 

class count_iterator(object):
  n = 0
 
  def __iter__(self):
    return self
 
  def next(self):
    y = self.n
    self.n += 1
    return y

下面是例子，注意最后一行试图将一个迭代器对象转为list，这将导致一个无限循环，因为这种迭代器对象将不会停止。
 

>>> counter = count_iterator()
>>> next(counter)
0
>>> next(counter)
1
>>> next(counter)
2
>>> next(counter)
3
>>> list(counter) # This will result in an infinite loop!

最后，我们将修改以上的程序：如果一个对象没有__iter__方法但定义了__getitem__方法，那么这个对象仍然是可迭代的。在这种情况下，当Python的内建函数iter将会返回一个对应此对象的迭代器类型，并使用__getitem__方法遍历list的所有元素。如果StopIteration或IndexError异常被抛出，则迭代停止。让我们看看以下的例子：
 

class SimpleList(object):
  def __init__(self, *items):
    self.items = items
 
  def __getitem__(self, i):
    return self.items[i]

用法在此：
 

>>> a = SimpleList(1, 2, 3)
>>> it = iter(a)
>>> next(it)
1
>>> next(it)
2
>>> next(it)
3
>>> next(it)
Traceback (most recent call last):
 File "<stdin>", line 1, in <module>
StopIteration

现在来看一个更有趣的例子：根据初始条件使用迭代器生成Hofstadter Q序列。Hofstadter在他的著作《G?del, Escher, Bach: An Eternal Golden Braid》中首次提到了这个嵌套的序列，并且自那时候开始关于证明这个序列对所有n都成立的问题就开始了。以下的代码使用一个迭代器来生成给定n的Hofstadter序列，定义如下：

Q(n)=Q(n-Q(n-1))+Q(n?Q(n?2))

给定一个初始条件，举个例子，qsequence([1, 1])将会生成H序列。我们使用StopIteration异常来指示序列不能够继续生成了，因为需要一个合法的下标索引来生成下一个元素。例如如果初始条件是[1,2]，那么序列生成将立即停止。
 

class qsequence(object):
  def __init__(self, s):
    self.s = s[:]
 
  def next(self):
    try:
      q = self.s[-self.s[-1]] + self.s[-self.s[-2]]
      self.s.append(q)
      return q
    except IndexError:
      raise StopIteration()
 
  def __iter__(self):
    return self
 
  def current_state(self):
    return self.s

用法在此：
 

>>> Q = qsequence([1, 1])
>>> next(Q)
2
>>> next(Q)
3
>>> [next(Q) for __ in xrange(10)]
[3, 4, 5, 5, 6, 6, 6, 8, 8, 8]

生成器(Generator)是一种用更简单的函数表达式定义的生成器。说的更具体一些，在生成器内部会用到yield表达式。生成器不会使用return返回值，而当需要时使用yield表达式返回结果。Python的内在机制能够帮助记住当前生成器的上下文，也就是当前的控制流和局部变量的值等。每次生成器被调用都适用yield返回迭代过程中的下一个值。__iter__方法是默认实现的，意味着任何能够使用迭代器的地方都能够使用生成器。下面这个例子实现的功能同上面迭代器的例子一样，不过代码更紧凑，可读性更强。
 

def count_generator():
  n = 0
  while True:
   yield n
   n += 1

来看看用法：

>>> counter = count_generator()
>>> counter
<generator object count_generator at 0x106bf1aa0>
>>> next(counter)
0
>>> next(counter)
1
>>> iter(counter)
<generator object count_generator at 0x106bf1aa0>
>>> iter(counter) is counter
True
>>> type(counter)
<type 'generator'>

现在让我们尝试用生成器来实现Hofstadter's Q队列。这个实现很简单，不过我们却不能实现前的类似于current_state那样的函数了。因为据我所知，不可能在外部直接访问生成器内部的变量状态，因此如current_state这样的函数就不可能实现了(虽然有诸如gi_frame.f_locals这样的数据结构可以做到，但是这毕竟是CPython的特殊实现，并不是这门语言的标准部分，所以并不推荐使用)。如果需要访问内部变量，一个可能的方法是通过yield返回所有的结果，我会把这个问题留作练习。
 

def hofstadter_generator(s):
  a = s[:]
  while True:
    try:
      q = a[-a[-1]] + a[-a[-2]]
      a.append(q)
      yield q
    except IndexError:
      return

请注意，在生成器迭代过程的结尾有一个简单的return语句，但并没有返回任何数据。从内部来说，这将抛出一个StopIteration异常。

下一个例子来自Groupon的面试题。在这里我们首先使用两个生成器来实现Bernoulli