python迭代器与生成器

总结自python cookbook。
本文主要介绍迭代器和生成器的常见用法，如有需要可以点击相应目录，书本完整内容可以搜索获取。

手动访问迭代器

用手动的方式调用迭代器可以让我们精确的把控程序的执行过程。手动访问需要使用next函数

def manual_iter():
    with open('/etc/passwd') as f:
        try:
            while True:
                line = next(f)  # 访问迭代器的下一个元素
                print(line, end='')
        except StopIteration:
            pass

委托迭代

当我们自己的类也想构建迭代器时可以使用委托迭代。Python的迭代协议要求定义一个__iter__()，将迭代请求委托到内部迭代器中。

class Node:
    def __init__(self, value):
        self._value = value
        self._children = []

    def __repr__(self):
        return 'Node({!r})'.format(self._value)

    def add_child(self, node):
        self._children.append(node)

    def __iter__(self):
        return iter(self._children) # 将迭代委托到列表中

# Example
if __name__ == '__main__':
    root = Node(0)
    child1 = Node(1)
    child2 = Node(2)
    root.add_child(child1)
    root.add_child(child2)
    # Outputs Node(1), Node(2)
    for ch in root:
        print(ch)

使用生成器来进行迭代

生成器可以直接应用迭代器，帮助我们实现一种新的迭代模式。使用yield语句就会定义一个生成器。对于生成器可以直接使用next函数。

def frange(start, stop, increment):
    x = start
    while x < stop:
        yield x
        x += increment
        
>>> for n in frange(0, 4, 0.5):   # 生成器为可迭代对象
...     print(n)
...
0
0.5
1.0
1.5
2.0
2.5
3.0
3.5

反向迭代

实现反向迭代可以使用内建的reversed()函数。对于自定义的类，实现反向迭代需要实现__reversed__()方法。

class Countdown:
    def __init__(self, start):
        self.start = start

    # 正向迭代时的迭代器
    def __iter__(self):
        n = self.start
        while n > 0:
            yield n
            n -= 1

    #反向迭代时的迭代器
    def __reversed__(self):
        n = 1
        while n <= self.start:
            yield n
            n += 1

for rr in reversed(Countdown(30)):  # 反向迭代函数
    print(rr)
for rr in Countdown(30):# 正向迭代函数
    print(rr)

对迭代器的切片

迭代器没办法使用简单的切片操作。但是迭代器的功能实在是强大，这是就可以使用itertools中的islice来对迭代器进行切片。

>>> def count(n):
...     while True:
...         yield n
...         n += 1
...
>>> c = count(0)
>>> c[10:20]
Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
TypeError: 'generator' object is not subscriptable

>>> # Now using islice()
>>> import itertools
>>> for x in itertools.islice(c, 10, 20):
...     print(x)
...
10
11
12
13
14
15
16
17
18
19
>>>

跳过可迭代对象的前一部分元素

这里跳过读取有两种方法，1.使用dropwhile 2.islice函数。dropwhile用于我们不知道需要跳过多少元素的情况。islice用于我们已经知道需要跳过多少元素。

# dropwhile例子
>>> from itertools import dropwhile
>>> with open('/etc/passwd') as f:
...     for line in dropwhile(lambda line: line.startswith('#'), f):  # 跳过所有注释行
...         print(line, end='')
...
nobody:*:-2:-2:Unprivileged User:/var/empty:/usr/bin/false
root:*:0:0:System Administrator:/var/root:/bin/sh
...
>>>

# isllice例子
>>> from itertools import islice
>>> items = ['a', 'b', 'c', 1, 4, 10, 15]
>>> for x in islice(items, 3, None): # 3为起始索引，这句话的意思是读取除前三个之外的所有元素
...     print(x)
...
1
4
10
15
>>>

利用迭代器获得排列组合

有时候我们需要获得一组元素的全排列，itertools自带了能够实现这种功能的工具。实现全排列利用itertools 中的 permutations。实现全部组合利用itertools中的combinations。

>>> items = ['a', 'b', 'c']
>>> from itertools import permutations
>>> for p in permutations(items):
...     print(p)
...
('a', 'b', 'c')
('a', 'c', 'b')
('b', 'a', 'c')
('b', 'c', 'a')
('c', 'a', 'b')
('c', 'b', 'a')
>>>
# 如果想要指定元素，可以添加相应的参数
>>> for p in permutations(items, 2):  # 获得前两个元素的全排列
...     print(p)
...
('a', 'b')
('a', 'c')
('b', 'a')
('b', 'c')
('c', 'a')
('c', 'b')
>>>

当我们需要的是元素的组合时，就可以用到 combinations函数。

>>> from itertools import combinations
>>> for c in combinations(items, 3):
...     print(c)
...
('a', 'b', 'c')

>>> for c in combinations(items, 2):
...     print(c)
...
('a', 'b')
('a', 'c')
('b', 'c')

>>> for c in combinations(items, 1):
...     print(c)
...
('a',)
('b',)
('c',) 

迭代时获得索引

如果我们需要在迭代的同时获得索引，那么可以使用enumerate函数。这样就不用为for循环无法获得元素索引而发愁了。

>>> my_list = ['a', 'b', 'c']
>>> for idx, val in enumerate(my_list):
...     print(idx, val)
...
0 a
1 b
2 c

同时迭代多个对象

有时我们需要对很多对象进行迭代，这个时候可以zip函数来实现。

>>> xpts = [1, 5, 4, 2, 10, 7]
>>> ypts = [101, 78, 37, 15, 62, 99]
>>> for x, y in zip(xpts, ypts):
...     print(x,y)
...
1 101
5 78
4 37
2 15
10 62
7 99
>>>

zip(a, b) 会生成一个可返回元组 (x, y) 的迭代器，其中x来自a，y来自b。 一旦其中某个序列到底结尾，迭代宣告结束。 **因此迭代长度跟参数中最短序列长度一致。**如果需要和最长的序列保持一致，那么就使用zip_longest（）函数。

在不同的容器中进行迭代

itertolls.chain()可以实现在多个容器中进行迭代。其原理是将多个可迭代对象作为输入然后输出一个迭代器。有了这个方法我们就不需要写很多for循环了。

>>> from itertools import chain
>>> a = [1, 2, 3, 4]
>>> b = ['x', 'y', 'z']
>>> for x in chain(a, b): # 将a,b进行了“链接”
... print(x)
...
1
2
3
4
x
y
z
>>>

将迭代对象扁平化

有时我们需要处理列表套列表的对象。这时候使用yield from来实现扁平化。yield from 可以递归的返回对象。其原理是将可迭代对象作为种子例程进行递归。

from collections import Iterable

def flatten(items, ignore_types=(str, bytes)):
    for x in items:
        if isinstance(x, Iterable) and not isinstance(x, ignore_types):
            yield from flatten(x) #返回子例程
        else:
            yield x

items = [1, 2, [3, 4, [5, 6], 7], 8]
# Produces 1 2 3 4 5 6 7 8
for x in flatten(items):
    print(x)

在上面代码中， isinstance(x, Iterable) 检查某个元素是否是可迭代的。 如果是的话， yield from 就会返回所有子例程的值。最终返回结果就是一个没有嵌套的简单序列了。

合并有序序列

合并有序序列可以使用heapq.merge()函数，其输入两个有序序列之后返回一个有序序列。heapq.merge()在合并的时候并不会将提供的序列一次性读取。

>>> import heapq
>>> a = [1, 4, 7, 10]
>>> b = [2, 5, 6, 11]
>>> for c in heapq.merge(a, b):
...     print(c)
...
1
2
4
5
6
7
10
11