Python 中的迭代

⚠ 转载请注明出处：作者：ZobinHuang，更新日期：Sept.23 2022

    本作品由 ZobinHuang 采用 知识共享署名-非商业性使用-禁止演绎 4.0 国际许可协议 进行许可，在进行使用或分享前请查看权限要求。若发现侵权行为，会采取法律手段维护作者正当合法权益，谢谢配合。

目录

有特定需要的内容直接跳转到相关章节查看即可。

正在加载目录...

前言

    在 Python 中我们经常会碰到和 Iteration (迭代) 相关的概念，比如我们使用 for 循环从一个列表中迭代元素就是一个经典的例子。殊不知，我们实际上还可以为我们自定义的类创建迭代方法，以使得它们也可以向列表一样被 for 循环等关键字迭代使用，经典的例子包括 PyTorch 中的 数据加载器 (Data Loader) 就运用了相关的技巧。本文中我们将理清与迭代相关的概念。

Magic Method (魔法方法)

    首先，我们需要知道: Python 抽象类拥有一些内置的可能会被各种关键字调用的 Magic Method (魔法方法)，常见的魔法方法以及它们对应的功能具体如下:

1. __len__(self): 定义当当前实体被 len() 函数调用时的行为，一般返回当前实体中包含的元素的个数;
2. __getitem__(self): 定义获取当前实体中指定元素时的行为，相当于 self[key]，即允许类对象拥有索引操作;
3. __iter__(self): 定义当迭代当前实体中的元素时的行为，一般行为会被定义为获取一个可用于访问当前实体元素的迭代器，Python 的 iter() 内置函数调⽤的就是对象的 __next__(self) 魔法⽅法;
4. __next__(self): 定义基于某个迭代器访问当前实体中下一个元素时的行为，Python 的 next() 内置函数调⽤的就是对象的 __next__(self) 魔法⽅法;

关于 Iterable 和 Iterator 概念的区分

    在 Python 中，Iterable 可以解释为可迭代的对象，其定义了魔法方法 __iter__，当该方法被触发时，它会返回一个 Iterator，也即迭代器。Iterator 顾名思义就是用于迭代访问 Iterable 元素的工具，其实现了 __iter__ 和 __next__ 两个魔法方法。Iterator 的 __iter__ 方法基本都是返回自身，而 __next__ 方法一般用于访问当前 Iterator 在 Iterable 中访问进度中的下一个元素，当访问到达末尾时则抛出 StopIteration 异常 iterator_iterable_diff。img_iterator_iterable 中说明了它们之间的关系。

    在 Python 程序中，我们一般首先对一个 Iterable 调用 iter() 函数以获得它的 Iterator，然后对新获得的 Iterator 调用 next() 函数以从 Iterable 中获取元素。

    使用 Python 的 for 等迭代关键字迭代某个 Iterable 时，背后会做的工作是:

1. 调用 Iterable 的 __iter__ 魔法方法以获得 Iterator;
2. 调用 Iterator 的 __next__ 魔法方法，企图获得下一个元素:
   • 若未抛出异常，则返回下一个获得的元素，for 循环取值返回;
   • 若抛出 StopIteration 异常，则迭代结束，for 循环结束;

    另外，我们有时会看见有一些程序的写法将 Iterable 和 Iterator 合二为一，举一个简单的例子如下所示:

class Colors:
  def __init__(self):
      self.rgb = ['red', 'green', 'blue']
      self.__index = 0

  def __iter__(self):
      return self

  def __next__(self):
      if self.__index >= len(self.rgb):
          raise StopIteration

      # return the next color
      color = self.rgb[self.__index]
      self.__index += 1
      return color

if __name__ == '__main__':
  colors = Colors()

  for color in colors:
      print(color)

Generator 和 yield 关键字

    我们在上面看到的 Iterable 是事先把数据元素创建好放置于内存中，然后再通过使用 Iterator 的方式对数据元素进行访问。而 Generator 是一种特殊的 Iterator+Iterable 结合体，它的数据元素并不是实现创建好的，它定义了一种数据创建的方式，然后在对这些数据元素发起迭代访问的时候采用 "on-the-fly" 的方式予以创建，一个简单的例子如下所示:

# 注意这里使用的是 () 而不是 []
mygenerator = (x*x for x in range(3))
for i in mygenerator:
  print(i)

    值得注意的是，Generator 有且只能被迭代访问一次。

    在完成对 Generator 的理解后，我们来看 yield 关键字。yield 功能类似于 return，不同的是 yield 返回的是一个 Generator。举例如下:

def create_generator():
    mylist = range(3)
    for i in mylist:
      yield i*i

# 创建一个 Generator
mygenerator = create_generator() 

# mygenerator 是一个 Generator 对象
print(mygenerator) 
# 打印结果: <generator object create_generator at 0xb7555c34>

# on-the-fly 创建数据
for i in mygenerator:
  print(i)

    在上面的代码中，我们把带有 yield 关键字的函数称为 Generator Function。当我们调用 Generator Function 的时候，实际上 Python 解释器并不会运行 Generator Function 函数体内的程序逻辑，而是会直接返回一个 Generator 对象。当我们第一次调用 for 循环对 Generator 发起迭代访问时，Generator Function 将会从头开始被执行，直到它运行至 yield 关键字处，并将 yield 关键字后面的值作为本次迭代获得的值进行返回; 在后续的对 Generator 的迭代访问中，都会从前一次 Generator Function 停止的下一行程序位置开始执行。当某一次对 Generator Function 的运行直到结束都没有碰到 yield 关键字时，则解释器判定迭代过程到底。

    我们有时还会看见 yield from 关键字，这个关键字允许 Generator Function 从另外一个 Generator Function 中返回 Generator 对象，如下所示:

def sample_generator(i):
  for j in range(i):
      yield j

def yf_generator(i):
  yield from sample_generator(i)

for value in yf_generator(5):
  print(value)