迭代器与生成器
迭代是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)