Python迭代器和生成器

Python中的迭代器是指我们可以对其进行迭代的对象。迭代器由可计数的值组成，并且可以一一遍历这些值。

迭代器仅实现Python的迭代器协议。迭代器协议是Python类，它带有两个特殊方法，即__iter __()和__next __()。使用这两种方法，迭代器可以计算迭代中的下一个值。

使用迭代器，我们很容易在Python中处理项目序列。我们不必将计算资源分配给序列中的所有项目，而是一次一次迭代单个项目，这有助于节省内存空间。

下面我们就一起来看一下Python迭代器和生成器相关知识。

基础知识

在本节中，将介绍和迭代器相关的一些基础知识。

容器

在Python当中一切皆对象，对象的抽象就是类，对象的集合就是容器，注意这里所说的集合不是通常数学里面的集合，而是一个广义的集合，这个里面是可以有重复元素的。

列表([1, 1, 2, 3])，元组((1, 1, 2, 3))，字典({0: 0, 1: 1})，集合({1, 2, 3})都是容器，对于容器，可以直观的理解成为多个元素组合在一起的单元，而不同容器的区别在于实现容器的内部的数据结构的实现方法，因此可以针对不同的场景来选择合适的容器。

可迭代对象

在这里，所有的容器都是可迭代的，注意这里的迭代不同于枚举，他们有一些小的差别。举个简单的栗子，比如自动售卖机，我们买东西的时候只需要告诉自动售卖机我要哪一个东西，大部分自动售卖机并不显示剩余的货物数量，每次要一个产品，自动售卖机就会吐出来一件商品，要是没有了，售卖机就会显示售空的状态，作为消费者，并不知道自动售卖机内部是如何存储这些货物的。

用稍微专业一点的话来说，迭代器提供了一个next方法，调用这个方法之后要么得到下一个容器的对象，要么得到一个StopInteraton的异常，这里不需要想列表那样指定元素的索引，实际上集合和字典压根就没有索引。

而迭代对象，通过iter函数返回一个迭代器，在通过next函数就可以实现遍历，for in iter_object，这实际上是一个语法糖，将调用next的过程隐式化了。

下面，我们来看看常见的数据类型是否可迭代。

def is_iterable(param):
    try: 
        iter(param) 
        return True
    except TypeError:
        return False

obj_list = [
    123,
    '123',
    [1, 2, 3],
    {1, 2, 3},
    {1: 1, 2: 2, 3: 3},
    (1, 2, 3)
]

for obj in obj_list:
    print(f'{obj} is iterable? {is_iterable(obj)}')

这里，除了数字类型之外，其他的数据类型都是可迭代的。

生成器

实际上，生成器可以看做是懒人版本的迭代器

在迭代器中，如果我们想枚举它的元素，这些元素要提前生成，来看下面一个例子:

import os
import psutil

def show_memory_info(hint):
    pid = os.getpid()
    p = psutil.Process(pid)
    
    info = p.memory_full_info()
    memory = info.uss / 1024. / 1024
    print('{} memory used: {} MB'.format(hint, memory))


def test_iterator():
    show_memory_info('initing iterator')
    list_1 = [i for i in range(100000)]
    show_memory_info('after iterator initiated')
    print(sum(list_1))
    show_memory_info('after sum called')

def test_generator():
    show_memory_info('initing generator')
    list_2 = (i for i in range(100000))
    show_memory_info('after generator initiated')
    print(sum(list_2))
    show_memory_info('after sum called')

%time test_iterator()
%time test_generator()

下面的%time是一个Jupyter的语法糖，可以输出函数运行的时间。

这里，我们可以发现使用列表的时候每个元素都会保存到内存当中，如果数据量巨大的话，将会占用海量的内存，如果内存不足极容易出现OOM的错误。

不过，我们并不需要在内存中同时保存这么多东西，比如对元素求和，我们只需要知道每个元素在相加的那一刻是多少就行了，用完就可以扔掉了。于是，生成器的概念应运而生，在你调用next函数的时候，才会生成下一个变量。

生成器在 Python 的写法是用小括号括起来，(i for i in range(10000000))，即初始化了一个生成器。这样一来，你可以清晰地看到，生成器并不会像迭代器一样占用大量内存，只有在被使用的时候才会调用。而且生成器在初始化的时候，并不需要运行一次生成操作。

生成器用法举例

下面，我们来具体看几个使用生成器的例子

一个简单的恒等式

这时候，我们首先来看一个小学二年级学过的一个恒等式 $(\sum_{i=1}^{n}{i})^2 = \sum_{i=1}^{n}{i^3}$ ，下面我们用代码来简单验证一下它的正确性:

def generator(k):
    i = 1
    while True:
        yield i ** k
        i += 1

gen_1 = generator(1)
gen_3 = generator(3)

print(gen_1)
print(gen_3)

def get_sum(n):
    s1, s3 = 0, 0

def get_sum(n):
    sum_1, sum_3 = 0, 0
    for i in range(n):
        next_1 = next(gen_1)
        next_3 = next(gen_3)
        print(f'next_1 = {next_1}, next_3 = {next_3}')
        sum_1 += next_1
        sum_3 += next_3
    print(sum_1 * sum_1, sum_3)

get_sum(8)

这里，看起来像是一个死循环，实际上也是，但是这个程序执行的效果并没有无限执行，实际上，这里生成器就发挥了巨大作用，这里的迭代器是一个无限集合，只要去调用next，就会无限生成下去，生成器就会根据运算自动生层新的元素，返回给调用者，灰常方便。

返回列表元素的索引

给定一个 list 和一个指定数字，求这个数字在 list 中的位置。

这个问题，常规做法比较简单，很容易写出如下的代码:

def get_index(l, target):
    result = []
    for i, num in enumerate(l):
        if num == target:
            result.append(i)
    return result

test_list = [2, 3, 2, 3, 2, 3, 1, 0, 5, 3]
print(f'2 in list of {test_list} is {get_index(test_list, 2)}')

再来看一下迭代器版:

def get_index(l, target):
    for i, num in enumerate(l):
        if num == target:
            yield i

test_list = [2, 3, 2, 3, 2, 3, 1, 0, 5, 3]
print(f'2 in list of {test_list} is {list(get_index(test_list, 2))}')

注意迭代器返回的是一个Generator对象，所以需要进行list转换之后才能使用print进行输出。

总结

本文介绍了四种不同的对象, 容器、可迭代对象、迭代器、生成器。

容器是可迭代对象，可迭代对象调用iter函数，可以得到一个迭代器。迭代器可以通过next函数来得到下一个元素，从而支持遍历。
生成器是一种特殊的迭代器（注意这个逻辑关系反之不成立）。使用生成器，你可以写出来更加清晰的代码；合理使用生成器，可以降低内存占用、优化程序结构、提高程序速度。

Python迭代器和生成器