进阶篇：Python迭代器深入理解

一、可迭代对象 Iterable

可迭代对象 Iterable：表示该对象可迭代，其并不是指某种具体数据类型。简单来说只要是实现了 __iter__ 方法的类就是可迭代对象。

from collections.abc import Iterable, Iterator

class A(object):
    def __init__(self):
        self.a = [1, 2, 3]

    def __iter__(self):
        # 此处返回啥无所谓
        return self.a

cls_a = A()
#  True
print(isinstance(cls_a, Iterable))

但是对象如果是 Iterable 的，看起来好像也没有特别大的用途，因为你依然无法迭代，实际上 Iterable 仅仅是提供了一种抽象规范接口：

for a in cls_a:
    print(a)

# 程序报错，要理解这个错误的含义
TypeError: iter() returned non-iterator of type 'list'

我们可以检查下 Iterable 接口：

class Iterable(metaclass=ABCMeta):

    # 如果实现了这个方法，那么就是 Iterable
    @abstractmethod
    def __iter__(self):
        while False:
            yield None

    @classmethod
    def __subclasshook__(cls, C):
        if cls is Iterable:
            return _check_methods(C, "__iter__")
        return NotImplemented

看起来实现 Iterable 接口用途不大，其实不是的，其有很多用途的，例如简化代码等，在后面的高级语法糖中会频繁用到，后面会分析。

二、迭代器 Iterator

迭代器 Iterator：其和 Iterable 之间是一个包含与被包含的关系，如果一个对象是迭代器 Iterator，那么这个对象肯定是可迭代 Iterable；但是反过来，如果一个对象是可迭代 Iterable，那么这个对象不一定是迭代器 Iterator，可以通过接口协议看出：

class Iterator(Iterable):

    # 迭代具体实现
    @abstractmethod
    def __next__(self):
        'Return the next item from the iterator. When exhausted, raise StopIteration'
        raise StopIteration

    # 返回自身，因为自身有 __next__ 方法(如果自身没有 __next__，那么返回自身没有意义)
    def __iter__(self):
        return self

    @classmethod
    def __subclasshook__(cls, C):
        if cls is Iterator:
            return _check_methods(C, '__iter__', '__next__')
        return NotImplemented

可以发现：实现了 __next__ 和 __iter__ 方法的类才能称为迭代器，就可以被 for 遍历了。到目前为止这句话是正确的，但是当你读到后面就知道这句话不严谨。

class A(object):
    def __init__(self):
        self.index = -1
        self.a = [1, 2, 3]

    # 必须要返回一个实现了 __next__ 方法的对象，否则后面无法 for 遍历
    # 因为本类自身实现了 __next__，所以通常都是返回 self 对象即可
    def __iter__(self):
        return self

    def __next__(self):
        self.index += 1
        if self.index < len(self.a):
            return self.a[self.index]
        else:
            # 抛异常，for 内部会自动捕获，表示迭代完成
            raise StopIteration("遍历完了")

cls_a = A()
print(isinstance(cls_a, Iterable)) # True
print(isinstance(cls_a, Iterator)) # True
print(isinstance(iter(cls_a), Iterator)) # True

for a in cls_a:
    print(a)
# 打印 1 2 3

再次明确，一个对象如果要是 Iterator ，那么必须要实现 __next__ 和 __iter__ 方法(不严谨，看到后面就知道为啥不严谨了)，但是要理解其内部迭代流程，还需要理解 for .. in .. 流程。

三、for .. in .. 本质流程

for .. in .. 也就是常见的迭代操作了，其被 python 编译器编译后，实际上代码是：

# 实际调用了 __iter__ 方法返回自身，包括了 __next__ 方法的对象
cls_a = iter(cls_a)
while True:
    try:
        # 然后调用对象的 __next__ 方法，不断返回元素
        value = next(cls_a)
        print(value)
    # 如果迭代完成，则捕获异常即可
    except StopIteration:
        break

可以看出，任何一个对象如果要能够被 for 遍历，必须要实现 __iter__ 和 __next__ 方法，缺一不可。同样的，这句话在我们已经解读过的知识下是正确的，但是不够严谨，甚至表述有点问题，因为该对象不需要直接实现这两个方法也可以得到一个迭代器，后面会说。严谨说法应该是该对象直接或者间接实现了这两个方法都可以返回迭代器，例如后面要讲的高级语法糖功能。

明白了上述流程，那么迭代器对象 A，我们可以采用如下方式进行遍历：

myiter = iter(cls_a)
print(next(myiter))
print(next(myiter))
print(next(myiter))
# 因为遍历完了，故此时会出现错误： StopIteration: 遍历完了
print(next(myiter))

我们再来思考 python 内置对象 list 为啥可以被迭代？

b=list([1,2,3])
print(isinstance(b, Iterable)) # True
print(isinstance(b, Iterator)) # False

可以发现 list 类型是可迭代对象，但是其不是迭代器(即 list 没有 __next__ 方法)，那为啥 for .. in .. 可以迭代呢？

原因是 list 内部的 __iter__ 方法内部返回了具备 __next__ 方法的类，或者说调用 iter() 后返回的对象本身就是一个迭代器，当然可以 for 循环了。

b=list([1,2,3])
print(dir(b)) # 可以发现其存在 __iter__ 方法，不存在 __next__

b=iter(b) # 调用 list 内部的 __iter__，返回了具备 __next__ 的对象
print(isinstance(b, Iterable)) # True
print(isinstance(b, Iterator)) # True
print(dir(b)) # 同时具备 __iter__ 和 __next__ 方法

基于上述理解我们可以对 A 类代码进行改造，使其更加简单：

class A(object):
    def __init__(self):
        self.a = [1, 2, 3]
    # 我们内部又调用了 list 对象的 __iter__ 方法，故此时返回的对象是迭代器对象
    def __iter__(self):
        return iter(self.a)

cls_a = A()
print(isinstance(cls_a, Iterable)) # True
print(isinstance(cls_a, Iterator)) # False

for a in cls_a:
    print(a)
# 输出： 1 2 3

此时我们就实现了仅仅实现 Iterable 规范接口，但是又具备了 for .. in .. 功能，代码是不是比最开始的实现简单很多？ 这种写法应用也非常广泛，因为其不需要自己再次实现 __next__ 方法。

如果你想理解的更加透彻，那么可以看下面例子：

# 仅仅实现 __iter__ 
class A(object):
    def __init__(self):
        self.b = B()

    def __iter__(self):
        return self.b

# 仅仅实现 __next__
class B(object):
    def __init__(self):
        self.index = -1
        self.a = [1, 2, 3]

    def __next__(self):
        self.index += 1
        if self.index < len(self.a):
            return self.a[self.index]
        else:
            # 内部会自动捕获，表示迭代完成
            raise StopIteration("遍历完了")


cls_a = A()
cls_b = B()
print(isinstance(cls_a, Iterable)) # True
print(isinstance(cls_a, Iterator)) # False
print(isinstance(cls_b, Iterable)) # False
print(isinstance(cls_b, Iterator)) # False

print(type(iter(cls_a))) # B 对象
print(isinstance(iter(cls_a), Iterator)) # False

for a in cls_a:
    print(a)

# 输出： 1 2 3

自此我们知道了：一个对象要能够被 for .. in .. 迭代，那么不管你是直接实现 __iter__ 和 __next__ 方法(对象必然是 Iterator)，还是只实现 __iter__(不是 Iterator)，但是内部间接返回了具备 __next__ 对象的类，都是可行的。

但是除了这两种实现，还有其他高级语法糖，可以进一步精简代码。

四、getitem理解

上面说过 for .. in .. 的本质就是调用对象的 __iter__ 和 __next__ 方法，但是有一种更加简单的写法，你通过仅仅实现 __getitem__ 方法就可以让对象实现迭代功能。实际上任何一个类，如果实现了__getitem__ 方法，那么当调用 iter(类实例) 时候会自动具备__iter__ 和 __next__方法，从而可迭代了。

通过下面例子可以看出，__getitem__ 实际上是属于 iter和next` 方法的高级封装，也就是我们常说的语法糖，只不过这个转化是通过编译器完成，内部自动转化，非常方便。

class A(object):
    def __init__(self):
        self.a = [1, 2, 3]

    def __getitem__(self, item):
        return self.a[item]

cls_a = A()
print(isinstance(cls_a, Iterable))  # False
print(isinstance(cls_a, Iterator))  # False
print(dir(cls_a))  # 仅仅具备 __getitem__ 方法

cls_a = iter(cls_a)
print(dir(cls_a))  # 具备 __iter__ 和 __next__ 方法

print(isinstance(cls_a, Iterable))  # True
print(isinstance(cls_a, Iterator))  # True

# 等价于 for .. in ..
while True:
    try:
        # 然后调用对象的 __next__ 方法，不断返回元素
        value = next(cls_a)
        print(value)
    # 如果迭代完成，则捕获异常即可
    except StopIteration:
        break

# 输出： 1 2 3

而且 __getitem__ 还可以通过索引直接访问元素，非常方便

a[0] # 1
a[4] # 错误，索引越界

如果你想该对象具备 list 等对象一样的长度属性，则只需要实现 __len__ 方法即可

class A(object):
    def __init__(self):
        self.a = [1, 2, 3]

    def __getitem__(self, item):
        return self.a[item]

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

cls_a = A()
print(len(cls_a)) # 3

到目前为止，我们已经知道了第一种高级语法糖实现迭代器功能，下面分析另一个更简单的可以直接作用于函数的语法糖。

五、yield 生成器

生成器是一个在行为上和迭代器非常类似的对象，二者功能上差不多，但是生成器更优雅，只需要用关键字 yield 来返回，作用于函数上叫生成器函数，函数被调用时会返回一个生成器对象，生成器本质就是迭代器，其最大特点是代码简洁。

def func():
    for a in [1, 2, 3]:
        yield a

cls_g = func()
print(isinstance(cls_g, Iterator))  # True
print(dir(cls_g))  # 自动具备 __iter__ 和 __next__ 方法

for a in cls_g:
    print(a)

# 输出: 1 2 3

# 一种更简单的写法是用 ()
cls_g = (i for i in [1,2,3])

直观感觉和 __getitem__ 一样，也是高级语法糖，但是比 __getitem__ 更加简单，更加好用。

使用 yield 函数与使用 return 函数，在执行时差别在于：包含 yield 的方法一般用于迭代，每次执行时遇到 yield 就返回 yield 后的结果，但内部会保留上次执行的状态，下次继续迭代时，会继续执行 yield 之后的代码，直到再次遇到 yield 后返回。生成器是懒加载模式，特别适合解决内存占用大的集合问题。假设创建一个包含10万个元素的列表，如果用 list 返回不仅占用很大的存储空间，如果我们仅仅需要访问前面几个元素，那后面绝大多数元素占用的空间都白白浪费了，这种场景就适合采用生成器，在迭代过程中推算出后续元素，而不需要一次性全部算出。

六、迭代器的高级用法

6.1 dropwhile

在实际的应用当中，比如记录的日志或者是代码等等，一般来说头部都会附上一段说明，或者用注释标注或者是用特殊的符号标记。这些信息是给用到数据的程序员看的，当我们通过代码获取数据的时候，显然是希望可以过滤掉这些信息的。

比如我们有一段数据，它的开头用#做了一些注释：

# This is a data for student
# Rows 100

xiaoming, 17, 99;
xiaoli, 18, 98;
...

常规操作当中，我们会创建一个打开文件的迭代器，我们通过遍历这个迭代器去获取文件当中的数据：

with open('xxxx.txt') as f:
  for line in f:
    print(line)

如果只是用来输出还好，如果我们需要加工文件当中的数据，那么头部的注释信息就会干扰我们代码的运行。我们当然可以手动加入一些判断，但是这会比较麻烦，代码也不够美观。针对这个问题，一个比较好的解决方案是dropwhile。

dropwhile是itemtools当中的一个函数，它可以接收一个我们自定义的过滤函数和迭代器重新生成一个新的迭代器，这个新的迭代器当中会过滤掉之前迭代器头部不符合我们要求的数据：

在刚才的例子当中我们想要过滤掉头部加了#注释的部分，我们可以这么操作：

from itertools import dropwhile
with open('xxxx.txt') as f:
  for line in dropwhile(lambda line: line.startswith('#'), f):
    print(line)

这样出来的结果就没有头部我们不需要的内容了。

当我们知道头部不符合情况的数据的格式的时候，可以使用dropwhile来规定过滤的格式。如果我们知道需要过滤的条数，则可以使用另外一个工具，叫做islice，它的本质是一个切片函数，就像是Python当中数组的切片功能一样，可以切出迭代器当中指定片段的数据。

举个例子：

from itertools import dropwhile
with open('xxxx.txt') as f:
  for line in islice(f, 3, None):
    print(line)

这样我们就会从第三行开始获取，之前的数据会被过滤掉。它其实就代表着数组当中[3: ]的切片操作。

6.2 迭代排列组合

我们都知道在C++当中有一个叫做next_permutation的函数，可以传入一个数组，返回下一个字典序的排列。在Python当中也有同样的功能，但是是以迭代器的形式使用的。

举个简单的例子，比如我们有a, b, c三个元素，我们希望求出它的所有排列：

items = ['a', 'b', 'c']
from itertools import permutations

for p in permutations(items):
  print(p)

permutations还支持多传一个参数，比如上述的排列当中我们希望只保留前两个元素，除了切片之外，我们只需要多传一个参数就好了，like this：

for p in permutations(items, 2):
  print(p)

除了排列之外，itertools当中还支持组合，用法还是一样，只是把函数名称换成是combinations而已：

from itertools import combindations
for c in combinations(items):
  print(c)

在一般的组合当中，一个元素一旦被选中那么它接下来就会从候选集当中移除，再也不会被选中。如果我们希望获得有放回的组合，我们可以再换一个函数，这个函数名称有点长，但是名字倒也直观叫做combinations_with_replacement。但既然是有放回的抽样，我们需要设定元素的数量，否则抽样可以无限进行下去。

for c in combinations_with_replacement(items, 3):
  print(c)

6.3 迭代合并后的序列

上一篇文章当中我们介绍了zip可以同时迭代多个迭代器，除此之外还有一种情况是我们需要把多个迭代器串起来迭代。比如系统的日志打在了多个文件当中，我们希望找出其中有error的日志来分析。这个时候，我们希望的不是同时读取多个迭代器，而是希望能够有办法将多个迭代器的内容串联起来。这个功能就是itertools当中的chain方法，它接受多个迭代器，当我们遍历的时候，会自动将多个迭代器的内容串联起来，我们可以无缝迭代。

举个例子：

from itertools import chain
nums = [1, 2, 3]
chars = ['a', 'b', 'c']

for i in chain(nums, chars):
  print(i)

这样我们会把nums和chars当中的内容一起输出出来，就好像从头到尾只执行了一个迭代器一样。

你可能会说我们不用chain也可以实现啊，我们可以这样：

for i in nums + chars:
  print(i)

的确，从结果上来看这样也是行得通的。但是如果我们分析一下内部执行的时候的中间变量，会发现当我们执行nums+chars的时候，实际上是先创建了一个新的临时list。然后在这个list当中存储nums和chars的数据，也就是说我们迭代的其实是这个新的list。这带来的结果是我们额外开辟了一段内存，并且花费了一些时间。如果我们使用chain，它并不会有这样的中间变量，完全是通过迭代器来执行的迭代，非常节省内存，这也是chain的优点。

6.4 归并迭代的内容

对于归并操作我们应该都不陌生，在之前的归并排序以及一些题解的文章当中我们见过很多次。同样，我们在使用工具合并多个迭代器内容的时候，如果迭代器当中的内容有序，我们也可以对多个迭代器当中的元素进行归并，而不再需要我们自己手动操作。

使用我们之前介绍的heapq的库可以非常轻松地做到这一点，我们一起来看一个例子：

a = [1, 3, 5]
b = [2, 4, 6]

import heapq

for c in heapq.merge(a, b):
  print(c)

执行之后，我们会得到[1, 2, 3, 4, 5, 6]的结果。也就是说通过heapq.merge操作，我们把多个有序的迭代器合并到了一起。当然我们也可以自己合并，但如果我们只是需要利用当中的数据的话，使用merge操作可以节省内存空间。

七、小结

• list set dict等内置对象都是容器 container 对象，容器是一种把多个元素组织在一起的数据结构，可以逐个迭代获取其中的元素。容器可以用 in 来判断容器中是否包含某个元素。大多数容器都是可迭代对象，可以使用某种方式访问容器中的每一个元素。
• 在迭代对象基础上，如果实现了 __next__ 方法则是迭代器对象，该对象在调用 next() 的时候返回下一个值，如果容器中没有更多元素了，则抛出 StopIteration 异常。
• 对于采用语法糖 __getitem__ 实现的迭代器对象，其本身实例既不是可迭代对象，更不是迭代器，但是其可以被 for in 迭代，原因是对该对象采用 iter(类实例) 操作后就会自动变成迭代器。
• 生成器是一种特殊迭代器，但是不需要像迭代器一样实现__iter__和__next__方法，只需要使用关键字 yield 就可以，生成器的构造可以通过生成器表达式 ()，或者对函数返回值加入 yield 关键字实现。
• 对于在类的 __iter__ 方法中采用语法糖 yield 实现的迭代器对象，其本身实例是可迭代对象，但不是迭代器，但是其可以被 for .. in .. 迭代，原因是对该对象采用 iter(类实例) 操作后就会自动变成迭代器