典型的迭代器模式作用很简单——遍历数据结构。不过,即便不是从集合中获取元素,而是获取序列中即时生成的下一个值时,也用得到这种基于方法的标准接口。例如,内置的 range 函数用于生成有穷整数等差数列(Arithmetic Progression,AP),itertools.count 函数用于生成无穷等差数列。
下一节会说明 itertools.count 函数,本节探讨如何生成不同数字类型的有穷等差数列。
下面我们在控制台中对稍后实现的 ArithmeticProgression 类做一些测试,如示例 14-10 所示。这里,构造方法的签名是 ArithmeticProgression(begin, step[, end])。range() 函数与这个 ArithmeticProgression 类的作用类似,不过签名是 range(start, stop[, step])。我选择使用不同的签名是因为,创建等差数列时必须指定公差(step),而末项(end)是可选的。我还把参数的名称由 start/stop 改成了 begin/end,以明确表明签名不同。在示例 14-10 里的每个测试中,我都调用了 list() 函数,用于查看生成的值。
示例 14-10 演示
ArithmeticProgression类的用法
>>> ap = ArithmeticProgression(0, 1, 3)
>>> list(ap)
[0, 1, 2]
>>> ap = ArithmeticProgression(1, .5, 3)
>>> list(ap)
[1.0, 1.5, 2.0, 2.5]
>>> ap = ArithmeticProgression(0, 1/3, 1)
>>> list(ap)
[0.0, 0.3333333333333333, 0.6666666666666666]
>>> from fractions import Fraction
>>> ap = ArithmeticProgression(0, Fraction(1, 3), 1)
>>> list(ap)
[Fraction(0, 1), Fraction(1, 3), Fraction(2, 3)]
>>> from decimal import Decimal
>>> ap = ArithmeticProgression(0, Decimal('.1'), .3)
>>> list(ap)
[Decimal('0.0'), Decimal('0.1'), Decimal('0.2')]
注意,在得到的等差数列中,数字的类型与 begin 或 step 的类型一致。如果需要,会根据 Python 算术运算的规则强制转换类型。在示例 14-10 中,有 int、float、Fraction 和 Decimal 数字组成的列表。
示例 14-11 列出的是 ArithmeticProgression 类的实现。
示例 14-11
ArithmeticProgression类
class ArithmeticProgression:
def __init__(self, begin, step, end=None): ➊
self.begin = begin
self.step = step
self.end = end # None -> 无穷数列
def __iter__(self):
result = type(self.begin + self.step)(self.begin) ➋
forever = self.end is None ➌
index = 0
while forever or result < self.end: ➍
yield result ➎
index += 1
result = self.begin + self.step * index ➏
❶ __init__ 方法需要两个参数:begin 和 step。end 是可选的,如果值是 None,那么生成的是无穷数列。
❷ 这一行把 self.begin 赋值给 result,不过会先强制转换成前面的加法算式得到的类型。10
10Python 2 内置了 coerce() 函数,不过 Python 3 没有内置。开发者觉得没必要内置,因为算术运算符会隐式应用数值强制转换规则。所以,为了让数列的首项与其他项的类型一样,我能想到最好的方式是,先做加法运算,然后使用计算结果的类型强制转换生成的结果。我在 Python 邮件列表中问了这个问题,Steven D'Aprano 给出了妙极的答复(https://mail.python.org/pipermail/python-list/2014-December/682651.html)。
❸ 为了提高可读性,我们创建了 forever 变量,如果 self.end 属性的值是 None,那么 forever 的值是 True,因此生成的是无穷数列。
❹ 这个循环要么一直执行下去,要么当 result 大于或等于 self.end 时结束。如果循环退出了,那么这个函数也随之退出。
❺ 生成当前的 result 值。
❻ 计算可能存在的下一个结果。这个值可能永远不会产出,因为 while 循环可能会终止。
在示例 14-11 中的最后一行,我没有直接使用 self.step 不断地增加 result,而是选择使用 index 变量,把 self.begin 与 self.step 和 index 的乘积相加,计算 result 的各个值,以此降低处理浮点数时累积效应致错的风险。
示例 14-11 中定义的 ArithmeticProgression 类能按预期那样使用。这是个简单的示例,说明了如何使用生成器函数实现特殊的 __iter__ 方法。然而,如果一个类只是为了构建生成器而去实现 __iter__ 方法,那还不如使用生成器函数。毕竟,生成器函数是制造生成器的工厂。
示例 14-12 中定义了一个名为 aritprog_gen 的生成器函数,作用与 ArithmeticProgression 类一样,只不过代码量更少。如果把 ArithmeticProgression 类换成 aritprog_gen 函数,示例 14-10 中的测试也都能通过。11
11本书源码仓库(https://github.com/fluentpython/example-code)中的 14-it-generator/ 目录里包含 doctest,以及一个 aritprog_runner.py 脚本,用于测试 aritprog*.py 脚本的所有版本。
示例 14-12
aritprog_gen生成器函数
def aritprog_gen(begin, step, end=None):
result = type(begin + step)(begin)
forever = end is None
index = 0
while forever or result < end:
yield result
index += 1
result = begin + step * index
示例 14-12 很棒,不过始终要记住,标准库中有许多现成的生成器。下一节会使用 itertools 模块实现,那个版本更棒。
itertools模块生成等差数列Python 3.4 中的 itertools 模块提供了 19 个生成器函数,结合起来使用能实现很多有趣的用法。
例如,itertools.count 函数返回的生成器能生成多个数。如果不传入参数,itertools.count 函数会生成从零开始的整数数列。不过,我们可以提供可选的 start 和 step 值,这样实现的作用与 aritprog_gen 函数十分相似:
>>> import itertools >>> gen = itertools.count(1, .5) >>> next(gen) 1 >>> next(gen) 1.5 >>> next(gen) 2.0 >>> next(gen) 2.5
然而,itertools.count 函数从不停止,因此,如果调用 list(count()),Python 会创建一个特别大的列表,超出可用内存,在调用失败之前,电脑会疯狂地运转。
不过,itertools.takewhile 函数则不同,它会生成一个使用另一个生成器的生成器,在指定的条件计算结果为 False 时停止。因此,可以把这两个函数结合在一起使用,编写下述代码:
>>> gen = itertools.takewhile(lambda n: n < 3, itertools.count(1, .5)) >>> list(gen) [1, 1.5, 2.0, 2.5]
示例 14-13 利用 takewhile 和 count 函数,写出的代码流畅而简短。
示例 14-13 aritprog_v3.py:与前面的
aritprog_gen函数作用相同
import itertools
def aritprog_gen(begin, step, end=None):
first = type(begin + step)(begin)
ap_gen = itertools.count(first, step)
if end is not None:
ap_gen = itertools.takewhile(lambda n: n < end, ap_gen)
return ap_gen
注意,示例 14-13 中的 aritprog_gen 不是生成器函数,因为定义体中没有 yield 关键字。但是它会返回一个生成器,因此它与其他生成器函数一样,也是生成器工厂函数。
示例 14-13 想表达的观点是,实现生成器时要知道标准库中有什么可用,否则很可能会重新发明轮子。鉴于此,下一节会介绍一些现成的生成器函数。