31. 谨慎地迭代函数所收到的参数 （effective-python）

谨慎地迭代函数所收到的参数

如果函数接受的参数是个包含许多对象的列表，那么这份列表可能要迭代多次。例如，我们要分析美国德克萨斯州的游客数量。原始数据保存在一份列表中，其中的每个元素表示每年有多少游客到这个城市旅游（单位是百万）。我们现在要统计每个城市的游客数占游客总数的百分比。
为了求出这份数据，我们编写一个归一化函数normalize，它先把列表里的所有元素加起来求出游客总数，然后，分别用每个城市的游客数除以游客总数计算出该城市在总数据中所占的百分比。

def normalize(numbers):
    # 第一次遍历列表numbers
    total = sum(numbers)
    result = []
    # 第二次遍历
    for value in numbers:
        percent = 100 * value / total
        result.append(percent)
    return result

把一份范例数据放到列表中传给这个函数，可以得到正确的结果。

visits = [15, 35, 80]
percentages = normalize(visits)
assert percentages == 100.0
print(percentages)

执行结果：

可能存在的问题：

• 若数据量很大，则会导致内存溢出

为了应付规模更大的数据，我们现在需要让程序能够从文件中读取信息，并假设德克萨斯州所有城市的游客数都放在这份文件中。我们尝试用生成器来实现，因为这样做可以让我们把同样的功能套用在其他数据上面，例如分析全世界（而不仅仅是德克萨斯一州）各城市的游客数。那些场合的数据量与内存用量可能会比现在大得多。

path = 'my_numbers.txt'
with open(path, 'w') as f:
    for i in (15, 35, 80):
        f.write('%d\n' % i)

def read_visits(data_path):
    with open(data_path) as f:
        for line in f:
            yield int(line)

然后我们对read_visits所返回的迭代器（其实为生成器，不过生成器是一种特殊的迭代器，下面统称迭代器）调用normalize函数，期望获得各城市游客占游客总数的百分比这一结果。

# 创建迭代器
it = read_visits(path)
percentages = normalize(it)
print(percentages)

执行结果：

可以发现，结果为空，那这是为什么呢？
出现这种状况的原因在于，迭代器只能产生一次结果。假如迭代器或生成器已经抛出过StopIteration异常，继续用它来构造列表或是像normalize那样对它做for循环，那它不会给出任何结果。

it = read_visits(path)
print(list(it))
print(list(it))  # 迭代器内元素已用完

执行结果：

在分析结果之前，我们先看一下迭代器的工作原理：

遍历迭代器i_t

for i in i_t:
    do_something_to(i)

等价于

# 其实是调用了i_t.__iter__()，该方法返回一个迭代器对象
fetch = iter(i_t)
while True:
    # 其实是调用了fetch.__next__()，来获得下一个元素
    try:
        data = next(fetch)
    except StopIteration:
        break
    do_something_to(data)

上面的等价结果其实还有一个可关注的点，那就是似乎直接调用fetch.__next__()，来获得下一个元素不也行吗？为什么还要先调用iter(i_t)呢？

原因就是得先判定这是一个迭代器，才能去迭代他，因此iter(i_t)是用来判断，我们迭代的对象是否是一个迭代器对象，因此如果迭代器 i_t 没有实现__iter__()方法，那它就不是迭代器，反之则是。

现在我们继续分析结果：

还有个因素也很令人迷惑，那就是：在已经把数据耗完的迭代器上继续迭代，程序居然不报错。这个过程等价于程序在第二次迭代时，继续调用next(fetch)，但是呢直接抛出了异常，并break掉 while 循环了，那为什么不报错呢？原因就是for循环、list构造器以及Python标准库的其他函数，都认为迭代器在正常过程的操作过程中抛出StopIteration异常是很自然的。它们没办法区分，这个迭代器本身就没有数据可以提供，还是虽然有数据，但现在已经耗尽了。

为了解决这个问题，我们可以把外界传入的迭代器特意遍历一整轮，从而将其中的内容复制到一份列表里。这样的话，以后就可以用这份列表来处理数据了，那时想迭代多少遍都可以。

个人觉得这没啥用，因为本身就是用迭代器去替换list，你现在还要用迭代一遍复制到list中，这不是在重复操作吗？

下面的函数与原来的 normalize 函数功能相同，但是它会把收到的那个 numbers 迭代器所能提供的数据明确地复制到 numbers_copy 列表里。

def normalize_copy(numbers):
    numbers_copy = list(numbers)  # Copy the iterator
    total = sum(numbers_copy)
    result = []
    for value in numbers_copy:
        percent = 100 * value / total
        result.append(percent)
    return result

现在这个函数就可以正确处理 read_visits 生成器函数所返回的 it 值了。

it = read_visits(path)
percentages = normalize_copy(it)
print(percentages)
assert sum(percentages) == 100.0

执行结果：

这个办法虽然能解决问题，但如果输入给它的迭代器所要提供的数据量非常庞大，那么有可能导致程序在复制迭代器时，因耗尽内存而崩溃。所以，把这个方案用在大规模的数据集合上面，会抵消掉 read_visits 的好处。当时编写那样一个函数，就是不想让程序把整套数据全都读取进来，而是可以一条一条地处理。为了应对大规模的数据，其中一个变通方案是让 normalize 函数接受另外一个函数（也就是下面的 get_iter），使它每次要用使用迭代器时，都去向那个函数索要。

其实就是原本传给 normalize 函数的参数是一个迭代器，现在传给它一个函数，这个函数可以生成迭代器，这样 normalize 函数需要几个迭代器，就找这个函数要几个。就不会存在只有一个迭代器，数据耗尽的情况了。

def normalize_func(get_iter):
    total = sum(get_iter())  # New iterator
    result = []
    for value in get_iter():  # New iterator
        percent = 100 * value / total
        result.append(percent)
    return result

使用 normalize_func 函数时，需要传入一条 lambda 表达式，让这个表达式去调用 read_visits 生成器函数。这样 normalize_func 每次向 get_iter 索要迭代器时，程序都会给出一个新的迭代器。说白了，就是有两个迭代器对象，这两个迭代器都去读取同一份文件。

# 每次调用 lambda 匿名函数，他就会帮助我们调用一次 read_visits(path)，生成一个迭代器对象
percentages = normalize_func(lambda: read_visits(path))
print(percentages)
assert sum(percentages) == 100.0

执行结果：

这样做虽然可行，但传入这么一个 lambda 表达式显得有点儿生硬。要想用更好的办法解决这个问题，可以新建一个容器类，让它实现迭代器协议（iterator protocol）。

Python 的 for 循环及相关的表达式，正是按照迭代器协议来遍历容器内容的。Python 执行 for x in foo 这样的语句时，实际上会调用 iter(foo)，也就是把 foo 传给内置的 iter 函数。这个函数会触发名为 foo.__iter__ 的特殊方法，该方法必须返回迭代器对象（这个对象本身要实现 __next__ 的特殊方法）。最后，Python 会用迭代器对象反复调用内置的 next 函数，直到数据耗尽为止（如果抛出 StopIteration 异常，就表示数据已经迭代完了）。

这个过程看听上去比较复杂，但总结起来，其实只需要让你的类在实现 __iter__ 方法时，按照生成器的方式来写就好。

我们的目的就是想得到一个可以创建多个迭代器对象的类，来替代原先使用 lambda 匿名函数的方式。

下面定义这样一种可迭代的容器类，用来读取文件中的游客数据。

class ReadVisits:
    def __init__(self, data_path):
        self.data_path = data_path

    def __iter__(self):
        with open(self.data_path) as f:
            for line in f:
                yield int(line)

我们来看一下这个 __iter__ 方法的本质，其实和之前的 read__visits 方法是基本一致的，只不过原来需要通过一个函数（lambda 匿名函数）来帮我们调用 read__visits 方法，从而不断生成新的迭代器对象，现在改用类 ReadVisits 来帮我们完成，相当于一种封装，更加优雅。

def __iter__(self):
    with open(self.data_path) as f:
        for line in f:
            yield int(line)
# 等价于    
def read_visits(data_path):
    with open(data_path) as f:
        for line in f:
            yield int(line)

之后，我们只需要把新的容器传给最早的那个 normalize 函数运行即可，函数本身的代码不需要修改。

# 实现迭代器协议的容器对象
visits = ReadVisits(path)
percentages = normalize(visits)
print(percentages)
assert sum(percentages) == 100

执行结果：

这样做为什么可行呢？因为 normalize 函数里面的 sum 会触发 ReadVisits.__iter__，让系统分配一个新的迭代器对象给它。接下来，normalize 通过 for 循环计算每项数据占总值的百分比时，又会触发 __iter__，于是系统会分配另一个迭代器对象。这些迭代器各自推进，其中一个迭代器把数据耗尽，并不会影响其他迭代器。所以，在每一个迭代器上面遍历，都可以分别看到一套完整的数据。这种方案的唯一缺点，就是多次读取输入数据。即每次创建一个迭代器对象，它都会重新读取一遍文件。

明白了 ReadVisits 这种容器的工作原理后，我们就可以在编写函数和方法时先确认一下，收到的应该是像 ReadVisits 这样的容器，而不是普通的迭代器。按照协议，如果将普通的迭代器传给内置的 iter 函数，那么函数会把迭代器本身返回给调用者。反之，如果传来的是容器类型，那么 iter 函数就会返回一个新的迭代器对象。我们可以借助这条规律，判断输入值是不是这种容器。如果不是，就抛出 TypeError 异常，因为我们没办法在那样的值上面重复遍历。

def normalize_defensive(numbers):
    # 判断是否是迭代器
    if iter(numbers) is numbers:  # An iterator -- bad!
        raise TypeError('Must supply a container')
    total = sum(numbers)
    result = []
    for value in numbers:
        percent = 100 * value / total
        result.append(percent)
    return result

还有一种写法也可以检测出这样的问题。collections.abc 内置模块里定义了名为 Iterator 的类，它用在 isinstance 函数中，可以判断自己收到的参数是不是这种实例。如果是，就抛出异常。

from collections.abc import Iterator

def normalize_defensive(numbers):
    if isinstance(numbers, Iterator):
        raise TypeError('Must supply a container')
    total = sum(numbers)
    result = []
    for value in numbers:
        percent = 100 * value / total
        result.append(percent)
    return result

这种自定义容器的方案，最适合既不想像 normalize_copy 函数一样，把迭代器的提供的这套数据全部复制一份，**同时又要多次遍历这套数据的情况。**该方案不仅支持自定义的容器，还支持系统自带的列表类型，因为这些全都属于遵循迭代器协议的可迭代容器。

visits = [15, 35, 80]
percentages = normalize_defensive(visits)
assert sum(percentages) == 100.0

visits = ReadVisits(path)
percentages = normalize_defensive(visits)
assert sum(percentages) == 100.0

如果输入的是普通迭代器而不是容器，那么 normalize_defensive 就会抛出异常。

visits = [15, 35, 80]
# 根据列表创建一个迭代器
it = iter(visits)
normalize_defensive(it)

执行结果：

这种思路也适合异步的迭代器。

要点：

• 函数和方法如果要把收到的参数遍历很多遍，那就必须特别小心。因为如果这些参数为迭代器，那么程序可能得不到预期的值，从而出现奇怪的效果。
• Python 的迭代器协议确定了容器与迭代器应该怎样跟内置的 iter 及 next 函数交互，以实现 for 循环与相关的表达式所表示的逻辑。
• 要想让自定义的容器类型可以迭代，只需要把 __iter__ 方法实现为生成器即可。
  （实现 __iter__、__next__ 的那种类属于自定义的迭代器，而不是容器）。
• 可以把值传给 iter 函数，检测它返回的是不是那个值本身。如果是，就说明这是个普通的迭代器，而不是一个可迭代的容器。另外，也可以用内置的 isintance 函数判断该值是不是 collections.abc.Iterator 类的实例。

总结：

自定义迭代器：实现了 __iter__ 、__next__ 方法的类

class Iterator:
    def __init__(self):
        ...

    def __iter__(self):
        return self

    def __next__(self):
        ...
        return ...

生成器：使用 yield 关键字的函数

def generator():
    ...
    yield ...

可迭代容器： __iter__ 方法使用了 yield 关键字的类

class container:
    def __init__(self):
        ...

    def __iter__(self):
        ...
        yield ...