第四章:迭代器与生成器

4.1 手动遍历迭代器

为了手动的遍历可迭代对象,使用 next() 函数并在代码中捕获 StopIteration 异常。

1
2
3
4
5
6
7
8
def manual_iter():
with open('/etc/passwd') as f:
try:
while True:
line = next(f)
print(line, end='')
except StopIteration:
pass

4.2 代理迭代

只需要定义一个 __iter__() 方法,将迭代操作代理到容器内部的对象上去。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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)


if __name__ == '__main__':
root = Node(0)
child1 = Node(1)
child2 = Node(2)
root.add_child(child1)
root.add_child(child2)
for ch in root:
print(ch)
# Node(1)
# Node(2)
  • Python 的迭代器协议需要 __iter__() 方法返回一个实现了__next__()方法的
    迭代器对象。
  • 如果你只是迭代其他容器的内容,称为代理迭代,你无须担心底层是怎样实现的。你所要做的只是传递迭代请求既可。

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

yield 语句

  • 一个生成器函数主要特征是它只会回应在迭代中使用到的 next 操作。
  • 一旦生成器函数返回退出,迭代终止。
  • 我们在迭代中通常使用的 for 语句会自动处理这些细节,所以你无需担心。

4.4 实现迭代器协议

Python 的迭代协议要求一个__iter__()方法返回一个特殊的迭代器对象,这个迭代器对象实现了__next__()方法并通过 StopIteration 异常标识迭代的完成。

4.5 反向迭代

能够反向迭代的要求(二选一)

  • 对象的大小可预先确定
  • 对象实现了__reversed__()方法

如果两者都不符合,那你必须先将对象转换为一个列表才行,

1
2
3
f = open('somefile')
for line in reversed(list(f)):
print(line, end='')

使用__reversed__()方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class CountDown(object):
def __init__(self, start):
self.start = start

def __iter__(self):
num = self.start
while num >= 0:
yield num
num -= 1

def __reversed__(self):
num = 0
while num <= self.start:
yield num
num += 1

c = CountDown(20)

for num in c:
print(num)
for num in reversed(c):
print(num)

定义一个反向迭代器可以使得代码非常的高效,因为它不再需要将数据填充到一个列表中然后再去反向迭代这个列表。

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

想让你的生成器暴露外部状态给用户

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
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
if 'python' in line:
for lineno, hline in lines.history:
print('{}:{}'.format(lineno, hline), end='')

# somefile.txt内容如下:
# hello world
# this is a test
# of iterating over lines with a history
# python is fun


# 打印如下:
# 2:this is a test
# 3:of iterating over lines with a history
# 4:python is fun

4.7 迭代器切片

函数 itertools.islice() 适用于在迭代器和生成器上做切片操作。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import itertools

def count(n):
while True:
yield n
n += 1

c = count(0)

for x in itertools.islice(c,10,15):
print(x)
# 10
# 11
# 12
# 13
# 14
  • 迭代器和生成器不能使用标准的切片操作,因为它们的长度事先我们并不知道
  • 函数 islice() 返回一个可以生成指定元素的迭代器,它通过遍历并丢弃直到切片开始索引位置的所有元素。然后才开始一个个的返回元素,并直到切片结束索引位置。
  • islice() 会消耗掉传入的迭代器中的数据。必须考虑到迭代器是不可逆的这个事实

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

itertools.dropwhile() 函数 : 只是跳过开始部分

  • 使用时,你给它传递一个函数对象和一个可迭代对象。它会返回一个迭代器对象,丢弃原有序列中直到函数返回 Flase 之前的所有元素,然后返回后面所有元素。
  • 简单来说,该函数起到过滤作用,具体使用就类似于filter
    ==满足条件的值都会丢弃直到有元素不满足为止==
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import itertools

def count(n):
num = 1
while num <= n:
yield num
num += 1

# 因为1%2==1,满足条件舍弃,2%2==0.不满足,所以之后的都保留
for num in itertools.dropwhile(lambda num:num % 2 == 1 ,count(4)):
print(num)
# 2
# 3
# 4

函数 dropwhile() 和 islice() 其实就是两个帮助函数,为的就是避免写出下面这种冗余代码:

1
2
3
4
5
6
7
8
9
10
11
12
with open('/etc/passwd') as f:
# Skip over initial comments
while True:
line = next(f, '')
if not line.startswith('#'):
break

# Process remaining lines
while line:
# Replace with useful processing
print(line, end='')
line = next(f, None)

4.9 排列组合的迭代

迭代遍历一个集合中元素的所有可能的排列或组合

itertools 模块提供了三个函数 :

  • itertools.permutations() : 排序
  • itertools.combinations() : 组合
  • itertools.combinations_with_replacement() : 组合,但是允许同一个元素被选择多次
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
from itertools import permutations

items = ['a','b','c']

# 排列
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')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from itertools import combinations

items = ['a','b','c']

# 排列
for p in combinations(items,3):
print(p)
# ('a', 'b', 'c')

# 在排列后,对元组切片
for p in combinations(items,2):
print(p)
# ('a', 'b')
# ('a', 'c')
# ('b', 'c')

for p in combinations(items,1):
print(p)
# ('a',)
# ('b',)
# ('c',)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from itertools import combinations_with_replacement

items = ['a','b','c']

# 排列
for p in combinations_with_replacement(items,3):
print(p)
# ('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')

当我们碰到看上去有些复杂的迭代问题时,最好可以先去看看 itertools 模块。

4.10 序列上索引值迭代

想在迭代一个序列的同时跟踪正在被处理的元素索引。

enumerate() 函数

1
2
3
4
5
6
7
8
9
10
11
# 简单统计单词数量
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 ,会发现它是一个字典 (准确来讲是一个defaultdict ),对于每个单词有一个 key ,每个 key 对应的值是一个由这个单词出现的行号组成的列表。

4.11 同时迭代多个序列

同时迭代多个序列,每次分别从一个序列中取一个元素。

  • zip() 函数
  • itertools.zip_longest() 函数
1
2
3
4
5
6
7
8
list1 = ['hyl','dsz','czj']
list2 = [98,97,96]

for item in zip(list1,list2):
print(item)
# ('hyl', 98)
# ('dsz', 97)
# ('czj', 96)

4.12 不同集合上元素的迭代

itertools.chain() 方法

4.13 创建数据处理管道

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
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


if __name__ == '__main__':

# Example 1
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)

# Example 2
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))

4.14 展开嵌套的序列

使用yieldyield from

1
2
3
4
5
6
7
8
9
10
11
12
from collections.abc import Iterable

def flatten(alist):
for ele in alist:
if isinstance(ele,Iterable):
yield from flatten(ele)
else:
yield ele

items = [1, 2, [3, 4, [5, 6], 7], 8]
print(list(flatten(items)))
# [1, 2, 3, 4, 5, 6, 7, 8]

yield from 返回所有子例程的值

当然也可以改成:

1
2
3
4
5
6
7
def flatten(items, ignore_types=(str, bytes)):
for x in items:
if isinstance(x, Iterable) and not isinstance(x, ignore_types):
for i in flatten(x):
yield i
else:
yield x

对比:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from collections.abc import Iterable


def flatten(alist):
res = []
for ele in alist:
if isinstance(ele,Iterable):
res.extend(flatten(ele))
else:
res.append(ele)
return res

items = [1, 2, [3, 4, [5, 6], 7], 8]
print(flatten(items))
# [1, 2, 3, 4, 5, 6, 7, 8]

注意:
口诀中的假设函数能用的意思是 :

  • 假设函数能用,考虑后续主函数该如何处理
  • 例如 : 方法二的flatten(ele)如果能用,那么就会返回[3, 4, 5, 6, 7, 8],接下来要考虑如何处理,这里就可以使用extend接受这个列表
  • 同理 : 方法一的flatten(ele)如果能用,那么他就会不断的yield出值,此时我们就可以yield from他了

4.15 顺序迭代合并后的排序迭代对象

你有一系列排序序列,想将它们合并后得到一个排序序列并在上面迭代遍历。

使用heapq.merge() 函数

1
2
3
4
5
6
7
8
import heapq

a = [1, 4, 7, 10]
b = [2, 5, 6, 11]

# 自动对两个元素进行排序
print(list(heapq.merge(a,b)))
# [1, 2, 4, 5, 6, 7, 10, 11]

heapq.merge有点类似于collections.ChainMap

1
2
3
4
5
6
7
from collections import ChainMap

a = {'name':'hyl'}
b = {'age':20}

print(ChainMap(a,b))
# ChainMap({'name': 'hyl'}, {'age': 20})
  • heapq.merge能使用堆排序,每次都最小值
  • collections.ChainMap是用于合并字典的,heapq.merge适用于合并生成器/迭代器的
  • 强调的是 heapq.merge() 需要所有输入序列必须是排过序的。

4.16 迭代器代替 while 无限循环

你在代码中使用 while 循环来迭代处理数据,因为它需要调用某个函数或者和一
般迭代模式不同的测试条件。能不能用迭代器来重写这个循环呢?

1
2
3
4
5
6
7
CHUNKSIZE = 8192
def reader(s):
while True:
data = s.recv(CHUNKSIZE)
if data == b'':
break
process_data(data)

使用迭代器重写whille循环:

1
2
3
def reader2(s):
for chunk in iter(lambda: s.recv(CHUNKSIZE), b''):
process_data(data)
  • iter 函数一个鲜为人知的特性是它接受一个可选的 callable 对象和一个标记 (结
    尾) 值作为输入参数。
  • 当以这种方式使用的时候,它会创建一个迭代器,这个迭代器会不断调用 callable 对象直到返回值和标记值相等为止