测试开发之Python核心笔记（15）：迭代器与生成器

15.1 可迭代对象Iterable

还记得for循环吗？for循环可以循环迭代处理字符串以及列表、元组、字典、集合这些容器类型，知道为什么这些数据类型可以被for迭代吗？因为这些对象都是可迭代对象。

判断是否是可迭代对象，可以isinstance(obj, Iterable)判断，输出True表示obj对象是可迭代的（iterable）。

15.2 迭代器iterator

通过迭代器，程序员可以迭代非序列类型，就是除了列表、元素、字典和集合之外类型。

迭代器通过next()方法获取对象中的下一个元素，可以把它看做一个数据流，我们不知道他的长度，一边使用一边计算下一个数据，是一种惰性计算。迭代器对象从集合的第一个元素开始访问，直到所有的元素被访问完结束。

迭代器的主要优点是节约内存（循环过程中，数据不用一次读入，在处理文件对象时特别有用，因为文件也是迭代器对象）、不依赖索引取值、实现惰性计算(需要时再取值计算)；

with open('java.txt') as f:for line in f:print(line)

这样每次读取一行处理一行，而不是一次性将整个文件读入，节约内存。

迭代器有个缺点，就是只能遍历一遍，不能遍历第二遍。因为遍历访问时，只能向前一个一个访问数据，直到访问到没有数据了，不能回头。也就是说迭代器只能往前不能后退。

l = ['a', 'n', 's', 'd']
l_iterator = iter(l)  # 通过iter函数得到一个迭代器
for i in l_iterator:print(i)for i in l_iterator:  # 经过第一遍迭代后，迭代器中已经没有数据了，或者说已经访问到尾部了。print(i)

15.2.1 通过 iter() 函数得到迭代器

tuple、list、dict、str虽然是Iterable，却不是迭代器Iterator。可以通过 iter() 函数返回一个迭代器。通过 next() 方法就可以实现遍历，调用next方法，要么得到这个容器的下一个对象，要么得到一个 StopIteration 的错误。Python的for循环迭代tuple、list本质上是先将他们转成迭代器Iterator，然后通过不断调用next()函数实现的。

l_iterator=iter([1,2,3,4,5])

15.2.1 通过类实现迭代器

在Python中有很多通过类实现的迭代器，比如reversed()，enumerate()，通过查看源码可以发现，这些类都实现了__next__方法和 __iter__ 方法。

如果你想让一个类对象可以被迭代，那么把一个类实现成一个迭代器，就是让类继承于Iterable，然后重写两个方法：__next__方法和 __iter__ 方法。

__iter__ 方法返回一个特殊的迭代器对象。

__next__ 会返回下一个迭代器对象。

例如，实现一个递减迭代器，对某个正整数，依次递减 1，直到 0。

from collections.abc import Iterableclass Decrease(Iterable):def __init__(self, init):self.init = initdef __iter__(self):  # 返回对象本身return selfdef __next__(self):while 0 < self.init:self.init -= 1return self.init # 返回下一个raise StopIteration  #  通过 raise 终断nextfor i in Decrease(6):  # 可以用for循环迭代这个类对象了print(i)

15.3 itertools 模块

这个模块实现了一系列快速、高效的迭代器。这些迭代器本身或组合都很有用。这一小节就来介绍一些非常好用的迭代器。

Help(itertools)可以看到内置的迭代器。如果使用Pycharm编辑器，按两下Shift，输入itertools，勾选上Include non-project items，在弹出的页面上选择itertools文件，就可以进入itertools.py文件。你会看到这样内容，这里列出来了itertools模块提供的所有迭代器：

"""
Functional tools for creating and using iterators.Infinite iterators:
count(start=0, step=1) --> start, start+step, start+2*step, ...
cycle(p) --> p0, p1, ... plast, p0, p1, ...
repeat(elem [,n]) --> elem, elem, elem, ... endlessly or up to n timesIterators terminating on the shortest input sequence:
accumulate(p[, func]) --> p0, p0+p1, p0+p1+p2
chain(p, q, ...) --> p0, p1, ... plast, q0, q1, ...
chain.from_iterable([p, q, ...]) --> p0, p1, ... plast, q0, q1, ...
compress(data, selectors) --> (d[0] if s[0]), (d[1] if s[1]), ...
dropwhile(pred, seq) --> seq[n], seq[n+1], starting when pred fails
groupby(iterable[, keyfunc]) --> sub-iterators grouped by value of keyfunc(v)
filterfalse(pred, seq) --> elements of seq where pred(elem) is False
islice(seq, [start,] stop [, step]) --> elements fromseq[start:stop:step]
starmap(fun, seq) --> fun(*seq[0]), fun(*seq[1]), ...
tee(it, n=2) --> (it1, it2 , ... itn) splits one iterator into n
takewhile(pred, seq) --> seq[0], seq[1], until pred fails
zip_longest(p, q, ...) --> (p[0], q[0]), (p[1], q[1]), ... Combinatoric generators:
product(p, q, ... [repeat=1]) --> cartesian product
permutations(p[, r])
combinations(p, r)
combinations_with_replacement(p, r)
"""

下面我们重点学习其中的几个迭代器。

15.3.1 拼接迭代器chain

可将多个可迭代对象变成为单个迭代器。在Pycharm的导航栏中，选择Navigate——>File Structure，就可以列出itertools.py模块中所有的类。

点击chain类，可以看到chain的说明：

class chain(object):"""chain(*iterables) --> chain objectReturn a chain object whose .__next__() method returns elements from thefirst iterable until it is exhausted, then elements from the nextiterable, until all of the iterables are exhausted."""

文档的意思就是说，创建一个迭代器，它首先返回第一个可迭代对象中所有元素，接着返回下一个可迭代对象中所有元素，直到耗尽所有可迭代对象中的元素。可将多个可迭代对象变成为单个迭代器。举几个例子，在Pycharm中编辑：

from itertools import chainchain_iterator = chain("ABCDef", "1234")  # 两个可迭代对象："ABCDef"和"1234"
print(list(chain_iterator))  # ['A', 'B', 'C', 'D', 'e', 'f', '1', '2', '3', '4']

使用 chain() 的一个常见场景是，当你想对不同的集合中所有元素执行某些操作的时候。比如，你想对两个列表的元素都计算平方和并形成一个新列表。可以这样做：

from itertools import chaindef square_multi_iterables(*iterables):  # 定义一个生成器chain_iterator = chain(*iterables)  # 利用chain生成迭代器for l in chain_iterator:  # 对迭代器里面的元素求进行平方yield l * l  # yield 返回if __name__ == '__main__':lst1 = [1, 2, 3]lst2 = [4, 5, 6]for i in square_multi_iterables(lst1, lst2):print(i)print(list(square_multi_iterables(lst2, lst1)))

这种解决方案要比使用两个单独的循环更加优雅！

15.3.2 累积迭代器accumulate

先看看源码中怎么说：

class accumulate(object):"""accumulate(iterable[, func]) --> accumulate objectReturn series of accumulated sums (or other binary function results)."""

将可迭代对象应用到func函数上，返回可迭代对象的累积迭代器。如果func没有提供，则返回可迭代对象累计和组成的迭代器。

看个例子说明一下：

from itertools import accumulatelst = [1, 2, 3, 4, 5, 6]
for i in accumulate(lst):print(i)
print(list(accumulate(lst)))  # 返回 [1, 3, 6, 10, 15, 21]

这个例子没有提供func函数，默认是对原序列累计求和。再来看看提供func函数的用法。

from itertools import accumulatelst = [1, 2, 3, 4, 5, 6]
for i in accumulate(lst, lambda x, y: x * y):print(i)
print(list(accumulate(lst, lambda x, y: x * y)))  # 返回 [1, 2, 6, 24, 120, 720]

15.3.3 排列组合迭代器

依然是先看看代码怎么说：

class combinations(object):"""combinations(iterable, r) --> combinations objectReturn successive r-length combinations of elements in the iterable.combinations(range(4), 3) --> (0,1,2), (0,1,3), (0,2,3), (1,2,3)"""

返回由输入iterable可迭代对象中元素组成长度为 r 的子序列。组合按照字典序返回。所以如果输入 iterable 是有序的，生成的组合元组也是有序的。

即使元素的值相同，不同位置的元素也被认为是不同的。如果元素各自不同，那么每个组合中没有重复元素。

举个例子看看，如何根据 “ABCD”, 输出 [‘AB’, ‘AC’, ‘AD’, ‘BC’, ‘BD’, ‘CD’]

def my_combinations(iterables, length):for i in combinations(iterables, length):  yield "".join(i)if __name__ == '__main__':lst=[]for element in my_combinations("ABCD", 2):  # 两个元素组成的排列组合lst.append(element)print(lst)  # ['AB', 'AC', 'AD', 'BC', 'BD', 'CD']

15.3.3 压缩迭代器compress

看看源码怎么说：

class compress(object):"""compress(data, selectors) --> iterator over selected dataReturn data elements corresponding to true selector elements.Forms a shorter iterator from selected data elements using theselectors to choose the data elements."""

创建一个迭代器，它返回 data 中经 selectors 真值测试为 True 的元素。迭代器在两者较短的长度处停止。

print(list(compress('ABCDEF', [1, 0, 1, 0, 1, 1])))  # 输出['A', 'C', 'E', 'F']

相当于是下面的这段代码：

def compress(data, selectors):return (d for d, s in zip(data, selectors) if s)

15.3.4 丢弃迭代器dropwhile

看看源码的描述：

class dropwhile(object):"""dropwhile(predicate, iterable) --> dropwhile objectDrop items from the iterable while predicate(item) is true.Afterwards, return every element until the iterable is exhausted."""

创建一个迭代器，如果 predicate 为true，迭代器丢弃这些元素，然后返回其他元素。迭代器在 predicate 首次为false之前不会产生任何输出，所以可能需要一定长度的启动时间。

举个例子：

from itertools import dropwhileprint(list(dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1])))  # 输出 [6, 4, 1]

大致相当于：

def my_dropwhile(predicate, iterable):iterable = iter(iterable)for x in iterable:if not predicate(x):yield x  # 返回第一个不满足predicate的值后退出这个循环breakfor x in iterable:  # 接着循环剩下的元素yield xprint(list(my_dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1])))  # 输出 [6, 4, 1]

更多的迭代器，可以参考itertools.py源码。

15.3 生成器generator（特殊的迭代器）

带 yield 的函数是生成器，而生成器也是一种迭代器。所以，生成器也有上面那些迭代器的特点。

通俗理解 yield，可结合函数的返回值关键字 return，yield 是一种特殊的 return。说是特殊的 return，是因为执行遇到 yield 时，立即返回，这是与 return 的相似之处。

不同之处在于：下次进入函数时直接到 yield 的下一个语句，而 return 后再进入函数，还是从函数体的第一行代码开始执行。

声明一个迭代器很简单，通过列表解析式[i for i in range(100000000)]就可以生成一个包含一亿个元素的列表。每个元素在生成后都会保存到内存中（这个过程很慢），而且占用了巨量的内存，内存不够的话就会出现 OOM 错误。不过，有的时候，我们并不需要提前在内存中保存这么多东西，比如对元素累加求和，我们只需要知道每个元素在相加的那一刻是多少就行了，用完就可以扔掉了。

生成器generator不会提前生成好所有的值放到内存中，只有调用next方法（生成器的方法）时，才会生成下一个变量，从而加快代码执行速度并节省内存。

得到生成器的方式有两种：

15.3.1 小括号定义生成器

将[]换成()，就可以简单地把列表解析式改成生成器generator，就得到一个生成器(i for i in range(100000000))。列表解析式最常用，但是尽量把列表解析式变成生成器，因为这样节省内存，节省内存的思想应该处处体现在代码里，这样才能体现水平。

15.3.2 通过yield关键字定义生成器

def fib(max):n, a, b = 0, 0, 1while n < max:yield ba, b = b, a + bn = n + 1return 'done'

在每次调用next()的时候执行，遇到yield语句暂停，返回yeild表达式后面的值，再次调用next时从yield语句处继续执行。

15.3.3 生成器的价值

生成器除了可以利用惰性计算来节约内存，还能提高代码的可读性。例如，求一段文本中，每个单词的的起始下标。不用生成器的方案：

def index_words(text):result = []  # 存放下标if text:result.append(0)for index, letter in enumerate(text, 1):  # 第二个参数是1，如果不写是什么情况if letter == " ":result.append(index)return resultif __name__ == '__main__':enumerate_desc = "The enumerate object yields pairs containing a count (from start, which defaults to zero) and a value yielded by the iterable argument."print(index_words(enumerate_desc))

如果使用生成器：

def index_words(text):if text:yield 0  # 第一次返回for index, letter in enumerate(text, 1): if letter == " ":yield index  # 每次调用返回if __name__ == '__main__':enumerate_desc = "The enumerate object yields pairs containing a count (from start, which defaults to zero) and a value yielded by the iterable argument."for index in index_words(enumerate_desc):print(index)print(list(index_words(enumerate_desc)))

可以使用生成器的代码更加清晰，代码行数更少。在不使用生成器时，对于每次结果，首先要执行一个append操作，最终才返回result。而使用生成器时候，直接yield index，少了列表操作带来的干扰，一眼就能看出，代码是要返回index。

kafka消费数据可以用yield生成器。还有pytest的fixture函数中也用到了，还有启动webdrvier也是用yield。

工作中一定要多多使用生成器。但是要注意，因为生成器也是迭代器，因此生成器只能遍历一次。

15.4 练习题

求列表中某一个元素的所有下标 index组成的列表

def index_generator(L, target):for i, num in enumerate(L):if num == target:yield iprint(list(index_generator([1, 6, 2, 4, 5, 2, 8, 6, 3, 2], 2)))########## 输出 ##########[2, 5, 9]

index_generator 会返回一个 Generator 对象，需要使用 list 转换为列表后，才能用 print 输出。

给定两个序列，判定第一个是不是第二个的子序列。

LeetCode 链接如下：https://leetcode.com/problems/is-subsequence/ 。序列就是列表，子序列则指的是，一个列表的元素在第二个列表中都按顺序出现，但是并不必挨在一起。举个例子，[1, 3, 5] 是 [1, 2, 3, 4, 5] 的子序列，[1, 4, 3] 则不是。

要解决这个问题，常规算法是贪心算法。维护两个指针指向两个列表的最开始，然后对第二个序列一路扫过去，如果某个数字和第一个指针指的一样，那么就把第一个指针前进一步。第一个指针移出第一个序列最后一个元素的时候，返回 True，否则返回 False。

不过本小节，使用生成器来实现。

def is_subsequence(a, b):b = iter(b)return all(i in b for i in a)print(is_subsequence([1, 3, 5], [1, 2, 3, 4, 5]))
print(is_subsequence([1, 4, 3], [1, 2, 3, 4, 5]))########## 输出 ##########True
False