迭代器与生成器

1. Abstract

Python 编程中，迭代器与生成器是重要的工具，也就是常见的：

for ... in ...:
	...

或者

next(it)

目前我学习到的有三种实现方案：

• __iter__(self) 和 __next__(self) 方法；
• __getitem__(self, index) 方法；
• 生成器函数；

2. __iter__(self) 和 __next__(self) 方法

关于这两个函数，我是看的博文《python之__iter__函数与__next__函数》，把代码搬过来：

class test():
	def __init__(self,data=1):
		self.data = data

	def __iter__(self):
		return self

	def __next__(self):
		if self.data > 5:
			raise StopIteration
		else:
			self.data+=1
			return self.data

for item in test(3):
	print(item)

##### output #####
4
5
6

for … in … 这个语句其实做了两件事。第一件事是获得一个可迭代器，即调用了 __iter__() 函数。第二件事是循环的过程，循环调用 __next__() 函数。

下面我们来看看到底怎么执行的，添加输出信息：

class Iter(object):
	def __init__(self, num):
		self.__num = num
		self.__idx = 0

	def __iter__(self):
		print('执行 __iter__()')
		self.__idx = 0
		return self

	def __next__(self):
		print('执行 __next__()')
		result = self.__idx
		self.__idx += 1
		if result >= self.__num:
			raise StopIteration
		return result

执行 for ... in ...：

for i in Iter(5):
	print(i)

输出：

执行 __iter__()
执行 __next__()
0
执行 __next__()
1
执行 __next__()
2
执行 __next__()
3
执行 __next__()
4
执行 __next__()

可见，执行 for ... in ... 时，会先调用一次 __iter__(self)，再迭代数执行次 __next__(self)。那这样的话，上面的博主写的迭代器就是有问题的：__iter__(self) 函数没有执行任何初始化操作。没猜错的话，如果再执行 for ... in ...，则不会有任何输出，因为迭代器保持了上次迭代后的状态：

t = test(3)
for item in t:
	print(item)
for item in t:
	print(item)

只会输出一个：

4
5
6

再看看我的，每次调用 __iter__(self) 时都会初始化 self.__idx = 0：

it = Iter(3)
for i in it:
	print(i)
for i in it:
	print(i)

执行 __iter__()
执行 __next__()
0
执行 __next__()
1
执行 __next__()
2
执行 __next__()
执行 __iter__()  # 第二次
执行 __next__()
0
执行 __next__()
1
执行 __next__()
2
执行 __next__()

iter() 和 next() 函数

好多形如 __func__(self) 的魔法方法都有对应的 func() 函数，__iter__(self) 和 __next__(self) 对应的函数是 iter() 和 next()。下面用 while 循环遍历迭代器：

it = iter(Iter(3))
while it:  # 本以为可以像 java 一样检查有没有下一个，结果发现无此类用法
	try:
		print(next(it))
	except StopIteration:
		break  # 故，只能用 break 结束循环了

##### output #####
# 和 for ... in ... 一样的输出
执行 __iter__()
执行 __next__()
0
执行 __next__()
1
执行 __next__()
2
执行 __next__()

调用 iter() 函数，会执行迭代器对象的 __iter__(self)，next() 会执行迭代器对象的 __next__(self)。所以，for ... in ... 背后的执行可能就是上述 while 循环的操作。

必须要有 __iter__(self)

上面博主说，不必有 __iter__(self)：

这仅仅使用了 __next__(self)，但是，循环遍历变成了 for i in range(3)，倘若执行 for i in test(3)，就会报错：

TypeError: 'Iter' object is not iterable  # 不可迭代

调用 iter(t)，一样会报这个错。

总结：迭代器的 __iter__(self) 很重要，迭代时会首先执行它，做一些遍历前的操作。

3. __get_item__(self, idx)

推荐博文《Python.__getitem__方法》：


第一次看到 __get_item__(self, idx) 是在 pytorch 的 torch.utils.data.Dataset，不过那时听说需要两个方法：__len__(self) 和 __getitem__(self, idx)。现在看来，__len__(self) 不要也可以，只不过，不能调用 len() 查看长度了，但这又很常用。

下面看一看到底发生了什么，修改 __getitem__(self, index)：

def __getitem__(self, index):
	print(index)  # 输出 index
	return self.animals_name[index]

输出：

0
dog
1
cat
2
fish
3  # 有 3

while 循环：

it = iter(animals)
while it:
	try:
		print(next(it))
	except StopIteration:
		break

##### output #####
0
dog
1
cat
2
fish
3

可见，和 __iter__(self) 和 __next__(self) 方法是完全一样的，可调 iter() 和 next() 函数，抛出的异常也是 StopIteration。

总结：要说区别，我觉得是：__getitem__(self, index) 更像是通过下标遍历列表；而 __iter__(self) 和 __next__(self) 方法适合更复杂的情况，比如遍历前需要一些操作，就可以放在 __iter__(self) 方法内。

4. 生成器函数

在博文《Python 上下文管理器》的第 5.2 节已经有一些解释，那里把它用作上下文管理器，现在把它当作迭代器。

先看的给出的例子：

import sys

def fibonacci(n):  # 生成器函数 - 斐波那契
	a, b, counter = 0, 1, 0
	while True:
		if counter > n:
			return
		yield a
		a, b = b, a + b
		counter += 1

f = fibonacci(10)  # f 是一个迭代器，由生成器返回生成

while True:
	try:
		print(next(f), end=" ")
	except StopIteration:
		sys.exit()

可以看到，是可以用 next(f) 迭代 fibonacci(10) 的。能不能用 for ... in ... 迭代呢？

for i in f:
	print(i, end=" ")

输出了：

0 1 1 2 3 5 8 13 21 34 55 

可见，它的确是迭代器。print(f) 会输出：

<generator object fibonacci at 0x0000022946108C10>

再检查 __iter__(self) 是否还在起作用，在 while 循环结束之后(sys.exit() 改成 break)，再调用 next(f) 会报错：

StopIteration

但如果是：

f = iter(f)  # 我们期望能从头再来
next(f)

结果并未像想象中的那样从头再来，而是报了错 StopIteration。就算是：

for i in f:
	print(i, end=" ")
f = iter(f)
next(f)

也是报错。可见，虽然生成器能用 for ... in ... 遍历，但它不是迭代器，f = iter(f) 虽然没有报错，但它没有从头再来。

总结：生成器不是迭代器，iter() 不能使它从头再来。