迭代器和生成器

可迭代对象、迭代器、生成器

概念和关系

可迭代对象(Iterable)

定义:实现了 __iter__() 方法的对象,或者可以被 iter() 函数作用的对象。

特点

  • 可以用 for x in obj: 来遍历。
  • 调用 iter(obj) 会返回一个 迭代器(iterator)

常见的可迭代对象

list, tuple, str, dict, set, range, file, 等等

示例

lst = [1, 2, 3]
it = iter(lst)  # 得到一个迭代器
迭代器(Iterator)

定义:实现了 __iter__()__next__() 方法的对象。

特点

  • __iter__() 返回自身(return self)。
  • __next__() 每次返回一个元素,直到抛出 StopIteration 异常。
  • 只能往前走,不能回头(一次性消费)。

示例

lst = [1, 2, 3]
it = iter(lst)  # it 是一个迭代器
print(next(it))  # 输出 1
print(next(it))  # 输出 2
生成器(Generator)

生成器是一种特殊的迭代器,有两种创建方式:

  1. 生成器函数(使用 yield
def my_gen():
    yield 1
    yield 2
    yield 3

g = my_gen()  # g 是一个生成器,也就是迭代器
  1. 生成器表达式(generator expression)
g = (x * x for x in range(3))  # g 是一个生成器

特点

  • 省内存,懒加载。
  • 遇到 yield 暂停执行,下一次 next() 从暂停点继续。
图解
           +-----------------------+
           |    Iterable           |   ← 可以用 for 循环遍历它
           |  (__iter__ method)    |
           +-----------------------+
                     |
                     |  iter()
                     ↓
           +-----------------------+
           |    Iterator           |   ← 可以手动调用 next()
           | (__iter__, __next__)  |
           +-----------------------+
                     ^
                     |
      +---------------------------+
      |     Generator             |   ← 特殊的 Iterator(自动实现了)
      |   (使用 yield 创建的)       |
      +---------------------------+

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

def frange(start, stop, increment):
    x = start
    while x < stop:
        yield x
        x += increment

为了使用这个函数, 可以用for循环迭代它或者使用其他接受一个可迭代对象的函数(比如 sum() , list() 等)。示例如下:

>>> 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
>>> list(frange(0, 1, 0.125))
[0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875]
>>>

一个函数中需要有一个 yield 语句即可将其转换为一个生成器。 跟普通函数不同的是,生成器只能用于迭代操作。

在对象上实现迭代

在一个对象上实现迭代最简单的方式是使用一个生成器函数。

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)

在上面代码中, __iter__() 方法只是简单的将迭代请求传递给内部的 _children 属性。

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

这里的 iter() 函数的使用简化了代码, iter(s) 只是简单的通过调用 s.__iter__() 方法来返回对应的迭代器对象, 就跟 len(s) 会调用 s.__len__() 原理是一样的。

depth_first() 方法简单直观。 它首先返回自己本身并迭代每一个子节点并 通过调用子节点的 depth_first() 方法(使用 yield from 语句)返回对应元素。

反向迭代

通过在自定义类上实现 __reversed__() 方法来实现反向迭代。比如:

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()


with open('somefile.txt') as f:
    lines = linehistory(f)
    for line in lines:
        if 'python' in line:
            for lineno, hline in lines.history:
                print('{}:{}'.format(lineno, hline), end='')

__iter__() 方法中定义生成器不会改变任何的算法逻辑。 由于它是类的一部分,所以允许定义各种属性和方法来供用户使用。

迭代器切片

函数 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
>>>

要着重强调的一点是 islice() 会消耗掉传入的迭代器中的数据,必须考虑到迭代器是不可逆的这个事实,所以如果需要之后再次访问这个迭代器的话,那么就得先将它里面的数据放入一个列表中。

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

itertools.dropwhile() 函数使用时,传递一个函数对象和一个可迭代对象。 它会返回一个迭代器对象,丢弃原有序列中直到函数返回Flase之前的所有元素,然后返回后面所有元素。

>>> with open('/etc/passwd') as f:
... for line in f:
...     print(line, end='')
...
##
# User Database
#
# Note that this file is consulted directly only when the system is running
# in single-user mode. At other times, this information is provided by
# Open Directory.
...
##
nobody:*:-2:-2:Unprivileged User:/var/empty:/usr/bin/false
root:*:0:0:System Administrator:/var/root:/bin/sh
...
>>>


>>> from itertools import dropwhile
>>> with open('/etc/passwd') as f:
...     for line in dropwhile(lambda line: not 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
...
>>>

如果已经明确知道了要跳过的元素的序号的话,那么可以使用 itertools.islice() 来代替。比如:

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

排列组合的迭代

itertools模块提供了三个函数来解决这类问题。 其中一个是 itertools.permutations() , 它接受一个集合并产生一个元组序列,每个元组由集合中所有元素的一个可能排列组成。 也就是说通过打乱集合中元素排列顺序生成一个元组,比如:

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

使用 itertools.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() 来讲,元素的顺序已经不重要了。 也就是说,组合 ('a', 'b') 跟 ('b', 'a') 其实是一样的(最终只会输出其中一个)。

在计算组合的时候,一旦元素被选取就会从候选中剔除掉(比如如果元素’a’已经被选取了,那么接下来就不会再考虑它了)。 而函数 itertools.combinations_with_replacement() 允许同一个元素被选择多次,比如:

>>> for c in combinations_with_replacement(items, 3):
...     print(c)
...
('a', 'a', 'a')
('a', 'a', 'b')
('a', 'a', 'c')
('a', 'b', 'b')
('a', 'b', 'c')
('a', 'c', 'c')
('b', 'b', 'b')
('b', 'b', 'c')
('b', 'c', 'c')
('c', 'c', 'c')
>>>

enumerate()迭代器

enumerate() 函数返回的是一个 enumerate 对象实例, 它是一个迭代器,返回连续的包含一个计数和一个值的元组, 元组中的值通过在传入序列上调用 next() 返回。

将一个文件中出现的单词映射到它出现的行号上去,可以很容易的利用enumerate() 来完成:

word_summary = defaultdict(list)

with open('myfile.txt', 'r') as f:
    lines = f.readlines()

for idx, line in enumerate(lines):
    # Create a list of words in current line
    words = [w.strip().lower() for w in line.split()]
    for word in words:
        word_summary[word].append(idx)

处理完文件后打印 word_summary ,对于每个单词有一个 key ,每个 key 对应的值是一个由这个单词出现的行号组成的列表。 如果某个单词在一行中出现过两次,那么这个行号也会出现两次, 同时也可以作为文本的一个简单统计。

zip()迭代器

zip(a, b) 会生成一个可返回元组 (x, y) 的迭代器,其中x来自a,y来自b。 一旦其中某个序列到底结尾,迭代宣告结束。 因此迭代长度跟参数中最短序列长度一致。

>>> a = [1, 2, 3]
>>> b = ['w', 'x', 'y', 'z']
>>> for i in zip(a,b):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
>>>

若要对齐长度,则可使用 itertools.zip_longest() 函数来代替。比如:

>>> from itertools import zip_longest
>>> for i in zip_longest(a,b):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
(None, 'z')

>>> for i in zip_longest(a, b, fillvalue=0):
...     print(i)
...
(1, 'w')
(2, 'x')
(3, 'y')
(0, 'z')
>>>

不同集合上元素的迭代

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
>>>

itertools.chain() 接受一个或多个可迭代对象作为输入参数。 然后创建一个迭代器,依次连续的返回每个可迭代对象中的元素。 这种方式要比先将序列合并再迭代要高效的多。

创建数据处理管道

foo/
    access-log-012007.gz
    access-log-022007.gz
    access-log-032007.gz
    ...
    access-log-012008
bar/
    access-log-092007.bz2
    ...
    access-log-022008

124.115.6.12 - - [10/Jul/2012:00:18:50 -0500] "GET /robots.txt ..." 200 71
210.212.209.67 - - [10/Jul/2012:00:18:51 -0500] "GET /ply/ ..." 200 11875
210.212.209.67 - - [10/Jul/2012:00:18:51 -0500] "GET /favicon.ico ..." 404 369
61.135.216.105 - - [10/Jul/2012:00:20:04 -0500] "GET /blog/atom.xml ..." 304 -
...

为了处理这些文件,定义一个由多个执行特定任务独立任务的简单生成器函数组成的容器。

import os
import fnmatch
import gzip
import bz2
import re

def gen_find(filepat, top):
    '''
    Find all filenames in a directory tree that match a shell wildcard pattern
    '''
    for path, dirlist, filelist in os.walk(top):
        for name in fnmatch.filter(filelist, filepat):
            yield os.path.join(path,name)

def gen_opener(filenames):
    '''
    Open a sequence of filenames one at a time producing a file object.
    The file is closed immediately when proceeding to the next iteration.
    '''
    for filename in filenames:
        if filename.endswith('.gz'):
            f = gzip.open(filename, 'rt')
        elif filename.endswith('.bz2'):
            f = bz2.open(filename, 'rt')
        else:
            f = open(filename, 'rt')
        yield f
        f.close()

def gen_concatenate(iterators):
    '''
    Chain a sequence of iterators together into a single sequence.
    '''
    for it in iterators:
        yield from it

def gen_grep(pattern, lines):
    '''
    Look for a regex pattern in a sequence of lines
    '''
    pat = re.compile(pattern)
    for line in lines:
        if pat.search(line):
            yield line

可以很容易的将这些函数连起来创建一个处理管道。 比如,为了查找包含单词python的所有日志行。

lognames = gen_find('access-log*', 'www')
files = gen_opener(lognames)
lines = gen_concatenate(files)
pylines = gen_grep('(?i)python', lines)
for line in pylines:
    print(line)

甚至可以在生成器表达式中包装数据。 比如,下面这个版本计算出传输的字节数并计算其总和。

lognames = gen_find('access-log*', 'www')
files = gen_opener(lognames)
lines = gen_concatenate(files)
pylines = gen_grep('(?i)python', lines)
bytecolumn = (line.rsplit(None,1)[1] for line in pylines)
bytes = (int(x) for x in bytecolumn if x != '-')
print('Total', sum(bytes))

展开嵌套的序列

可以写一个包含 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)

语句 yield from 在想在生成器中调用其他生成器作为子例程的时候非常有用。

最后要注意的一点是, yield from 在涉及到基于协程和生成器的并发编程中扮演着更加重要的角色。

迭代器代替while无限循环

一个常见的IO操作程序可能会想下面这样:

CHUNKSIZE = 8192

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

这种代码通常可以使用 iter() 来代替,如下所示:

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

iter 函数一个鲜为人知的特性是它接受一个可选的 callable 对象和一个标记(结尾)值作为输入参数。 当以这种方式使用的时候,它会创建一个迭代器, 这个迭代器会不断调用 callable 对象直到返回值和标记值相等为止。

这种特殊的方法对于一些特定的会被重复调用的函数很有效果,比如涉及到I/O调用的函数。 举例来讲,如果你想从套接字或文件中以数据块的方式读取数据,通常你得要不断重复的执行 read()recv() , 并在后面紧跟一个文件结尾测试来决定是否终止。这节中的方案使用一个简单的 iter() 调用就可以将两者结合起来了。 其中 lambda 函数参数是为了创建一个无参的 callable 对象,并为 recvread() 方法提供了 size 参数。

单例模式

工厂模式