python：列表生成式、生成器、迭代器

本文介绍python中常用的功能：列表生成式、生成器、迭代器。
主要参考：https://www.liaoxuefeng.com/wiki/1016959663602400/1017317609699776和https://www.runoob.com/python3/python3-iterator-generator.html

1 列表生成式

列表生成式即List Comprehensions，是Python内置的非常简单却强大的可以用来创建list的生成式。

1.1 最简单的列表生成式

假如要生成list: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]，可以用list(range(1, 11))：

>>> list(range(1, 11))
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

1.2 for循环用于列表生成式

假如要生成[1x1, 2x2, 3x3, ..., 10x10]怎么办？

最容易想到的肯定是使用循环：

>>> L = []
>>> for x in range(1, 11):
...     L.append(x * x)
... 
>>> L
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

列表生成器可以用一行代码代替上面繁琐的循环：

>>> [ x * x for x in range(1, 11)]
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

1.3 if判断用于列表生成式

假如要生成仅包含偶数的平方的列表：[2x2, 4x4, ..., 10x10]，可在for循环后再加上if判断来筛选出偶数：

>>> [ x * x for x in range(1, 11) if x % 2 == 0]
[4, 16, 36, 64, 100]

1.4 两层循环的列表生成式

直接看例子好了：

>>> [ m + n for m in 'ABC' for n in 'XYZ']
['AX', 'AY', 'AZ', 'BX', 'BY', 'BZ', 'CX', 'CY', 'CZ']

类似的，可以在列表生成式中使用三层及更多层的循环，只不过用到的比较少。

1.5 使用多个变量迭代的for循环用于列表生成式

>>> a = ['Hello', 'world', 'python', 'pytorch']
>>> ['%d: %s' % (i, s) for i, s in enumerate(a)]
['0: Hello', '1: world', '2: python', '3: pytorch']

2 生成器

通过列表生成式，我们可以直接创建一个列表。但是，受到内存限制，列表容量肯定是有限的。而且，创建一个包含100万个元素的列表，不仅占用很大的存储空间，如果我们仅仅需要访问前面几个元素，那后面绝大多数元素占用的空间都白白浪费了。

所以，如果列表元素可以按照某种算法推算出来，那我们是否可以在循环的过程中不断推算出后续的元素呢？这样就不必创建完整的list，从而节省大量的空间。在Python中，这种一边循环一边计算的机制，称为生成器：generator。

2.1 generator简单的创建方法

最简单的创建generator的方法是：把一个列表生成式的[]改为()，就创建了一个generator：

>>> L = [ x * x for x in range(1, 11)]
>>> L
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
>>> g = ( x * x for x in range(1, 11))
>>> g
<generator object <genexpr> at 0x7f1473558fc0>

2.2 访问generator的返回值

可以通过next()函数一个一个访问generator的返回值：

>>> g = ( x * x for x in range(1, 11))
>>> next(g)
1
>>> next(g)
4
>>> next(g)
9
>>> next(g)
16
>>> next(g)
25
>>> next(g)
36
>>> next(g)
49
>>> next(g)
64
>>> next(g)
81
>>> next(g)
100
>>> next(g)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

我们讲过，generator保存的是算法，每次调用next(g)，就计算出g的下一个元素的值，直到计算到最后一个元素，没有更多的元素时，抛出StopIteration的错误。

更方便的方法是使用for循环访问generator的返回值：

>>> g = ( x * x for x in range(1, 11))
>>> for n in g:
...     print(n)
... 
1
4
9
16
25
36
49
64
81
100

有一点需要注意，如果我们通过next()访问了一部分元素，再通过for循环访问的时候，这部分不会在被访问：

>>> g = ( x * x for x in range(1, 11))
>>> next(g)
1
>>> next(g)
4
>>> for n in g:
...     print(n)
... 
9
16
25
36
49
64
81
100

2.3 使用yield创建generator

有一些generator的推算算法比较复杂，使用类似列表生成式的for循环无法实现，那就需要使用函数来实现。

例如著名的斐波拉契数列（Fibonacci），除第一个和第二个数外，其他数字都由前两个数相加得到：
1, 1, 2, 3, 5, 8, 13, 21, 34, ...
使用列表生成式写不出来，但是用函数可以很方便的把它打印出来：

>>> def fib(max):
...     n, a, b = 0, 0, 1
...     while n < max:
...             print(b)
...             a, b = b, a + b
...             n = n + 1
...     return 'done'
... 
>>> fib(6)
1
1
2
3
5
8
'done'

如果一个函数的定义在包含yield关键字，那么这个函数就不再是一个普通函数，而是一个generator。
下面我们使用yield关键字将fib()函数改写为一个generator：

>>> def fib(max):
...     n, a, b = 0, 0, 1
...     while n < max:
...             yield b
...             a, b = b, a + b
...             n = n + 1
...     return 'done'
... 
>>> f = fib(6)
>>> print(f)
<generator object fib at 0x7f14735880a0>

这里，最难理解的就是generator和函数的执行流程不一样。函数是顺序执行，遇到return语句或者最后一行函数语句就返回。而变成generator的函数，在每次调用next()的时候执行，遇到yield语句返回，再次执行时从上次返回的yield语句处继续执行。

举个简单例子，定义一个generator，一次返回数字1, 3, 5：

>>> def odd():
...     print('step 1')
...     yield 1
...     print('step 2')
...     yield 3
...     print('step 3')
...     yield 5
... 
>>> o = odd()
>>> next(o)
step 1
1
>>> next(o)
step 2
3
>>> next(o)
step 3
5
>>> next(o)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

可以看到，我们使用yield关键字将函数odd()变成一个generator。在执行过程中，遇到yield就中断，下次从中断的地方继续执行。

再回来看fib()的例子，我们使用for循环来获得下一个返回值：

>>> f = fib(6)
>>> for n in f:
...     print(n)
... 
1
1
2
3
5
8

但是我们这样做没办法拿到fib()的返回值return ’done'。如果想要拿到fib()的返回值，需要捕获StopIteration异常，而返回值就包含在StopIteration的value中：

>>> f = fib(6)
>>> while True:
...     try:
...             print('f: %d' % next(f))
...     except StopIteration as e:
...             print('Generator return value: %s' % e.value)
...             break
... 
f: 1
f: 1
f: 2
f: 3
f: 5
f: 8
Generator return value: done

3 迭代器

1. 可直接用于for循环的数据类型有两类：

   1. 集合数据类型，如list、tuple、dict、set、str等；
   2. 生成器generator。

   这些可以直接作用于for循环的对象统称为可迭代对象：Iterable。
   实际上，只要是内置有__iter__方法的对象，都称为Iterable。

2. 可以被next()函数调用并不断返回下一个值的对象称为迭代器：Iterator，所以generator是迭代器。
   实际上，只要是内置有__next__方法的对象，都称为Iterator。

3. 迭代器有两个基本方法：iter()和next()。

4. Iterable和Iterator是两个不同的概念，实际上，Iterable包含了Iterator。

3.1 从集合数据类型创建Iterator

最简单的创建迭代器的方法是使用iter()函数从集合数据类型如list、tuple等创建：

>>> from collections import Iterator, Iterable
>>> list = [1, 2, 3, 4]
>>> isinstance(list, Iterable)
True
>>> isinstance(list, Iterator)
False
>>> isinstance(iter(list), Iterator)
True

3.2 自定义一个迭代器

把一个类作为一个迭代器使用需要在类中实现两个方法：__iter__()和__next__()。

class MyNumbers():
    def __init__(self, max_num=1):
        self.a = 0
        self.max_num = max_num

    def __iter__(self):
        return self

    def __next__(self):
        if self.a >= self.max_num:
            raise StopIteration
        else:
            self.a += 1
        return self.a

test = MyNumbers(8)
for n in test:
    print(n)

输出为：

1
2
3
4
5
6
7
8

• 执行第16行时，会自动调用MyNumbers类的__init__方法。
• 执行第17行的for循环时其实做了两件事：
  • 首先调用MyNumbers类的__iter__方法获取可迭代对象；
  • 循环过程，即循环调用__next__方法。

而在使用next()函数时，用不到__iter__方法，下面我们改写这个程序，注释掉__iter__方法：

class MyNumbers():
    def __init__(self, max_num=1):
        self.a = 0
        self.max_num = max_num

    #  def __iter__(self):
    #      return self

    def __next__(self):
        if self.a >= self.max_num:
            raise StopIteration
        else:
            self.a += 1
        return self.a

#  test = MyNumbers(8)
#  for n in test:
#      print(n)

test = MyNumbers(3)
print(next(test))
print(next(test))
print(next(test))
print(next(test))

输出为：

1
2
3
Traceback (most recent call last):
  File "test.py", line 24, in <module>
    print(next(test))
  File "test.py", line 11, in __next__
    raise StopIteration
StopIteration

shell returned 1

可以看到前面三个next()函数都成功执行，直到迭代完所有元素触发StopIteration异常。
这是如果再使用第17~18行的for循环，则会报错：

Traceback (most recent call last):
  File "test.py", line 17, in <module>
    for n in test:
TypeError: 'MyNumbers' object is not iterable

shell returned 1

提示MyNumbers类是不可迭代的。