PythonCookBook笔记——迭代器与生成器

迭代器与生成器

迭代是Python最强大的功能之一，虽然看起来迭代只是处理序列中元素的一种方法，但不仅仅如此。

手动遍历迭代器

想遍历但不想使用for循环。

使用next()方法并在代码中捕获StopIteration异常。

StopIteration用来指示迭代的结尾，也可以通过返回指定结尾。

l = next(iterator, None)

代理迭代

构建了一个自定义容器对象，想在这个容器上执行迭代操作。

只需定义__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)

迭代器协议需要__iter__()方法返回一个实现了__next__()方法的迭代器对象。

用生成器创建新的迭代模式

只需要实现yield语句即可转换为生成器，并且生成器只能用于迭代操作。

实现迭代器协议

最简单的是使用生成器函数，否则需要实现__iter__()和__next__()方法并完成对StopIteration异常的捕捉。

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)

    def depth_first(self):
        yield self
        for c in self:
            yield from c.depth_first()

# Example
if __name__ == '__main__':
    root = Node(0)
    child1 = Node(1)
    child2 = Node(2)
    root.add_child(child1)
    root.add_child(child2)
    child1.add_child(Node(3))
    child1.add_child(Node(4))
    child2.add_child(Node(5))

    for ch in root.depth_first():
        print(ch)
    # Outputs Node(0), Node(1), Node(3), Node(4), Node(2), Node(5)

class Node2:
    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)

    def depth_first(self):
        return DepthFirstIterator(self)

class DepthFirstIterator(object):
    '''
    Depth-first traversal
    '''

    def __init__(self, start_node):
        self._node = start_node
        self._children_iter = None
        self._child_iter = None

    def __iter__(self):
        return self

    def __next__(self):
        # Return myself if just started; create an iterator for children
        if self._children_iter is None:
            self._children_iter = iter(self._node)
            return self._node
        # If processing a child, return its next item
        elif self._child_iter:
            try:
                nextchild = next(self._child_iter)
                return nextchild
            except StopIteration:
                self._child_iter = None
                return next(self)
        # Advance to the next child and start its iteration
        else:
            self._child_iter = next(self._children_iter).depth_first()
            return next(self)

通常没人会去写第二种复杂的代码，又要维护状态又要处理异常，因此最好是使用生成器来实现。

反向迭代

使用reversed()方法，并且要注意，反向迭代必须是对象大小确定或该对象实现了__reversed__()方法才能生效，两者都不符合就需要将对象转换为列表才行。

# Print a file backwards
f = open('somefile')
for line in reversed(list(f)):
    print(line, end='')

还要注意的是，如果可迭代对象元素很多，转换为列表会消耗大量内存。

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

    # Forward iterator
    def __iter__(self):
        n = self.start
        while n > 0:
            yield n
            n -= 1

    # Reverse iterator
    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)

带有外部状态的生成器函数

如果想定义一个生成器函数，同时调用某个想暴露给用户使用的状态值，最简单的方法是实现一个类，然后把生成器函数放到__iter__()方法中。

from collections import deque

class linehistory:
    def __init__(self, lines, histlen=3):
        self.lines = lines
        self.history = deque(maxlen=histlen)

    def __iter__(self):
        for lineno, line in enumerate(self.lines, 1):
            self.history.append((lineno, line))
            yield line

    def clear(self):
        self.history.clear()

迭代器切片

想得到一个由迭代器/生成器生成的切片对象，使用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
>>>

跳过可迭代对象的开始部分

itertools模块的dropwhile()函数，传入一个函数对象和一个可迭代对象，返回一个可迭代对象，类似filter()方法，丢弃函数返回True的元素。

>>> from itertools import dropwhile
>>> with open('/etc/passwd') as f:
...     for line in dropwhile(lambda line: line.startswith('#'), f):
...         print(line, end='')

如果知道元素个数，也可以用islice()方法来抛弃前n个元素。

>>> from itertools import islice
>>> items = ['a', 'b', 'c', 1, 4, 10, 15]
>>> for x in islice(items, 3, None):
...     print(x)
...
1
4
10
15
>>>

None的作用与切片的[3:]原理相同。

排列组合的迭代

有时需要遍历一个集合中元素的所有可能的排列或组合。

itertools模块提供了三个函数来解决此类问题。

permutaions()接受可迭代对象和可选的长度参数，生成基于指定长度的排列元组。

>>> 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',)
>>>

若要允许同一元素被多次选择，可使用combinations_with_replacement()方法。

枚举迭代

在迭代的同时跟踪被处理的元素的下标索引。

使用内置的enumerate()方法，接收一个可迭代对象和可选的初始值，返回一个迭代器。

同时迭代多个序列

使用zip()方法，该方法需要注意以最短序列长度为迭代基准，超过长度不迭代。

或可使用itertools.zip_longest()方法，接收多个可迭代对象和一个fillvalue关键字参数指定默认值，此方法会迭代到最长序列。

注意的是zip()方法返回一个迭代器而不是列表。

同时对多个可迭代对象进行迭代

避免写重复的循环，使用itertools模块的chain()方法组合多个可迭代对象并返回一个新的迭代器。

>>> from itertools import chain
>>> a = [1, 2, 3, 4]
>>> b = ['x', 'y', 'z']
>>> for x in chain(a, b):
... print(x)
...
1
2
3
4
x
y
z
>>>

数据处理管道

使用生成器函数来实现管道机制。读取文件做一个生成器，读取行做一个生成器，处理行做一个生成器，最后用一个循环或相应方法调用，形成一个数据管道，注意的是yield和yield from的区别。

展开嵌套的序列

利用yield from后接可迭代对象会返回其所有元素的特点来调用，避免重复的循环代码，更优雅。

注意要判断是否是可迭代对象。yield from对于在生成器中调用其他生成器很有用。

合并有序序列并生成有序可迭代对象

有时需要将多个有序序列合并成一个有序序列，使用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

由于可迭代特性，heapq.merge()不会立刻读取所有序列，因此在长序列中使用不会有太大开销，并且必须注意的是，输入的序列必须是排序过的，heapq.merge()方法不会检查顺序，这个方法只是比较多个序列中的首位值，较小的放入新的序列。

迭代器代替while无限循环

利用了iter()方法的一个特性来做无限循环或有条件的循环。该方法接收一个可调用对象，iter()方法不断调用该对象直到其返回值与标记值相等为止。

>>> bool(iter(int, 1))
True

int默认值是0,因此iter迭代器永远不会结束，所以其布尔值始终是True。

CHUNKSIZE = 8192

def reader(s):
    while True:
        data = s.recv(CHUNKSIZE)
        if data == b'':
            break
        process_data(data)

def reader2(s):
    for chunk in iter(lambda: s.recv(CHUNKSIZE), b''):
        pass
        # process_data(data)