Python核心丨迭代器和生成器

迭代器和生成器

---

可迭代对象和迭代器

在Python中一切皆对象，对象的抽象就是类，而对象的集合就是容器。

列表（list），元组（tuple），字典（dict），集合（set）都是容器。所有的容器都是可迭代的（iterable）。

迭代器提供了一个next的方法。调用这个方法后，要么得到这个容器的下一个对象，要么得到一个StopIteration的错误。

而可迭代对象，通过iter()函数返回一个迭代器，再通过next()函数就可以实现遍历。

判断一个对象是否可迭代

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

params = [
    1234,
    '1234',
    [1, 2, 3, 4],
    set([1, 2, 3, 4]),
    {1:1, 2:2, 3:3, 4:4},
    (1, 2, 3, 4)
]
    
for param in params:
    print('{} is iterable? {}'.format(param, is_iterable(param)))

########## 输出 ##########

1234 is iterable? False
1234 is iterable? True
[1, 2, 3, 4] is iterable? True
{1, 2, 3, 4} is iterable? True
{1: 1, 2: 2, 3: 3, 4: 4} is iterable? True
(1, 2, 3, 4) is iterable? True

使用isinstance也可以

生成器

生成器是懒人版本的迭代器

import os
import psutil

# 显示当前 python 程序占用的内存大小
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(100000000)]
    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(100000000))
    show_memory_info('after generator initiated')
    print(sum(list_2))
    show_memory_info('after sum called')

%time test_iterator()
%time test_generator()

########## 输出 ##########

initing iterator memory used: 48.9765625 MB
after iterator initiated memory used: 3920.30078125 MB
4999999950000000
after sum called memory used: 3920.3046875 MB
Wall time: 17 s
initing generator memory used: 50.359375 MB
after generator initiated memory used: 50.359375 MB
4999999950000000
after sum called memory used: 50.109375 MB
Wall time: 12.5 s

声明一个迭代器很简单，[i for i in range(100000000)]就可以生成包含一亿元素的列表。每个元素再生成后都会保存到内存中。

不过，并不需要再内存中同时保存这么多东西，比如对元素求和，只需要知道每个元素在相加的那一刻是多少就行了，用完就可以扔掉了。

于是，生成器的概念应运而生，再调用next()函数的时候，才会生成下一个变量。

生成器再Python的写法是用小括号括起来，(i for i in range(100000000)，即初始化了一个生成器。

生成器更多的作用

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):
    sum_1, sum_3 = 0, 0
    for i in range(n):
        next_1 = next(gen_1)
        next_3 = next(gen_3)
        print('next_1 = {}, next_3 = {}'.format(next_1, next_3))
        sum_1 += next_1
        sum_3 += next_3
    print(sum_1 * sum_1, sum_3)

get_sum(8)

########## 输出 ##########

<generator object generator at 0x000001E70651C4F8>
<generator object generator at 0x000001E70651C390>
next_1 = 1, next_3 = 1
next_1 = 2, next_3 = 8
next_1 = 3, next_3 = 27
next_1 = 4, next_3 = 64
next_1 = 5, next_3 = 125
next_1 = 6, next_3 = 216
next_1 = 7, next_3 = 343
next_1 = 8, next_3 = 512
1296 1296

generator()这个函数，它返回了一个生成器

yield是魔术的关键。可以理解为，函数运行到这一行的时候，程序会从这里暂停，然后跳出到next()函数。

每次next(gen)函数被调用的时候，暂停的程序就又复活了，从yield这里向下继续执行；同时局部变量 i 并没有被清除掉，而是会继续累加。

迭代器是一个有限集合，生成器则可以成为一个无限集。

实例

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

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

print(index_normal([1, 6, 2, 4, 5, 2, 8, 6, 3, 2], 2))

########## 输出 ##########

[2, 5, 9]

使用迭代器

def index_generator(L, target):
    for i, num in enumerate(L):
        if num == target:
            yield i

print(list(index_generator([1, 6, 2, 4, 5, 2, 8, 6, 3, 2], 2)))

########## 输出 ##########

[2, 5, 9]

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