vlambda博客
学习文章列表

在本文中，我们将尽可能以最好的方式重新实现 python 的内置函数：enumerate。

前言

我们将从一个粗略的、不能真正解决问题的实现方案开始，一点一点地对其进行修改，直到达成完全重新实现内置 enumerate 函数的目标。我们会以实现该目标为动力，研究一系列有趣的 python 概念及其微妙之处。在阅读本文的同时，您将会更好地理解 enumerate 是什么以及它是如何工作的。

阅读本文后，你将：

• 理解内置函数 enumerate 是什么及其是如何工作的
• 能够指出通用的可迭代对象和列表之间的差异
• 了解什么是（惰性）生成器
• 了解 zip和 enumerate 的关系
• 学会 yield 和 yield from  如何使用
• 了解 yield 和 yield from  的差异
• 了解“无限”生成器
• 在 python 中完全重新实现 enumerate 函数

序言

假设你有一个字符串：

s = "hey"

你该如何打印该字符串的每个字符？典型的初学者会这样写一个循环：

for idx in range(len(s)):
    print(s[idx])

h
e
y

不过初学者很快就了解到可迭代对象的强大功能，这使他们能够以更明确和更有表现力的方式编写前面的循环：

for char in s:
    print(char)

h
e
y

注意我们是如何为字符串的每个元素命名的（char），而不是必须使用整数 idx 来索引字符串。这是因为字符串是可迭代对象：一个可以遍历的、可以在循环中逐一获取子元素的对象。

假设现在你想打印每个字符及其索引值，而不是仅仅打印字符。你会怎样做呢？

也许你会持怀疑的态度，恢复使用 len(s)：

s = "rod"
for idx in range(len(s)):
    print(f"Letter {idx} is {s[idx]}")

Letter 0 is r
Letter 1 is o
Letter 2 is d

当然， python 提供了一个内置函数，让我们以富有表现力的方式在编写 for 循环的同时做到这一点：即 enumerate 函数。通过使用它，我们可以访问连续的元素及其索引：

s = "rod"
for idx, letter in enumerate(s):
    print(f"Letter {idx} is {letter}")

Letter 0 is r
Letter 1 is o
Letter 2 is d

在本文中，我们将研究内置的 enumerate 函数并尝试以自己的方式实现它。我们会从一个简单的枚举模型开始，修正它的行为并提高它的代码质量。

第一次尝试

在编写我们对 enumerate 的第一个重新实现的方案之前，查看 enumarate 产生的内容或许是一个好主意。当然，我们知道它提供了索引和元素，因为这就是我们在 for 循环中可以得到的...但是这些索引和元素从何而来？

for element in enumerate(s):    
    print(element)

(0, 'r')
(1, 'o')
(2, 'd')

正如所见， enumerate  产生了包含两个元素的元组：

1、索引

2、可迭代对象的对应元素

值得注意的是， enumerate(s) 也是一个可迭代对象，我们遍历了它并逐一打印出了元素。

鉴于  enumerate(s) 的元素似乎是元组，或许我们可以从这里起手。那么什么样的“简单”对象可以替换下面的 ... 并给出相同的输出呢？

for element in ...:  
    print(element)

(0, 'r')
(1, 'o')
(2, 'd')

一种可能性是包含了所有元组的列表：

for element in [(0, 'r'), (1, 'o'), (2, 'd')]:  
    print(element)

(0, 'r')
(1, 'o')
(2, 'd')

那么，也许我们可以通过返回元组列表的方式实现  enumerate ：

>>> enumerate_("rod")
[(0, 'r'), (1, 'o'), (2, 'd')]

为此，我们可以使用一个 for 循环，同时使用这些元组填充列表，然后返回即可：

def enumerate_(iterable):   
    result = []   
    idx = 0  
    for elem in iterable:  
        result.append((idx, elem))  
        idx += 1 
    return result

这是一个对我们有帮助的  enumerate  实现模型，因为通过它我们可以了解  enumerate  的动态输出。但它并不是一个非常准确的模型，为了理解其中缘由，我们需要讨论一下生成器。

惰性生成器

快速思考一下，哪种代码会输出以下的列表呢？

>>> ...
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

如果你想到了 range(10)，那么已经很接近了，但不是正确答案。让我们看看 range(10) 的真正输出：

>>> range(10)
range(0, 10)

range 本身并不能产生我们想要的数字，为什么呢？因为 range 是一个生成器。

生成器是一种知道如何生成一系列值，但是仅在需要目标值的时候才会生成元素的 python 对象。例如，range(10) 确切地知道如何生成数字 0 到 10，但是上面的代码，还没有真正调用到目标元素，所以并不会显示出对应的数字。如果想查看这些数字的话，可以使用 list 函数：

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

list 之所以能有这个效果，是因为 list 需要使用 range 产生的值来构建列表。

巧合的是， enumerate 也会返回这种（惰性）生成器：

>>> enumerate("rod")
<enumerate object at 0x000001F616DE4540>

如果想要查看特定的  enumerate  对象所产生的元素，也只需要调用 list 函数即可：

>>> list(enumerate("rod"))
[(0, 'r'), (1, 'o'), (2, 'd')]

生成器的出色之处在于每次仅生成出一个元素，这也是它被描述为“惰性”的原因。在很多不需要获取全部元素的场景，生成器可以节省大量的时间，因为它不需要进行预先计算。

那么我们该如何重构  enumerate_ 函数以实现一次返回一个元素？如何将“惰性”嵌入  enumerate_ 函数呢？我们希望可以实现的像这样：

def enumerate_(iterable): 
    idx = 0   
    for elem in iterable:     
        return idx, elem  
        idx += 1

当然，我们知道 return 会直接返回并结束函数，上述写法将永远返回 (0, elem)。此处的 elem 就是可迭代对象的第一个元素：

>>> enumerate_("rod")
(0, 'r')

我们所期待的 "return" 并不该是这样，而是应该在给出第一个元素之后，能够继续一个接一个地产生下一个元素。python 原本的 return 是一个具有特定含义的关键字，所以也许我们可以尝试使用另一个关键字，代码将会变成这样：

def enumerate_(iterable):  
    idx = 0  
    for elem in iterable:     
        <lazy-result-keyword> idx, elem   
        idx += 1

lazy-result-keyword 代表的含义就是暂时返回当前值，后续可以有更多的值返回。

这个关键字就是 yield：

def enumerate_(iterable): 
    idx = 0  
    for elem in iterable:   
        yield idx, elem     
        idx += 1

通过使用 yiled，我们的  enumerate_ 函数就实现了“惰性”，它不会一次返回所有的元素。可以通过这样调用来验证一下：

>>> enumerate_("rod")
<generator object enumerate_ at 0x000001BAB7372D60>

与内置的  enumerate  函数类似，我们可以调用 list 函数来获取所有的元素：

>>> list(enumerate_("rod"))
[(0, 'r'), (1, 'o'), (2, 'd')]

到目前为止，我们的  enumerate_ 可以像内置的  enumerate  一样返回生成器，但是这样的实现还不够准确。

可选参数 start

下一步需要做的是确保  enumerate_ 具有内置  enumerate  函数的所有功能。既然如此，我们接下来要实现尚未研究的可以选参数 start。

如上文提及的，当我们提供给  enumerate  一个可迭代对象时，它将用索引和元素构建元组，并且索引从 0 开始：

>>> list(enumerate("rod"))
[(0, 'r'), (1, 'o'), (2, 'd')]

不过， enumerate  可以接受第二个参数，该参数可以指定索引的计数起点：

>>> list(enumerate("rod", 5))
[(5, 'r'), (6, 'o'), (7, 'd')]

第二个参数即 start。

为了实现这个功能，我们需要在  enumerate_ 函数也添加第二个参数。然后，使用传入的参数初始化起始索引的值：

def enumerate_(iterable, start=0):  
    idx = start  
    for elem in iterable:    
        yield idx, elem    
        idx += 1

让我们试试：

>>> list(enumerate_("rod"))
[(5, 'r'), (6, 'o'), (7, 'd')]

到此为止，我们似乎用一段简单的代码实现了  enumerate  的所有功能，好像目标已经达成了。但真的是这样吗？让我们看看是否可以改进它。

记录索引

当我们开始着手于重构一段代码时，有一些地方很值得注意，它与代码算法本身没有太大关联，但却是整个工作所必需的。对于我们的  enumerate_ 函数，变量 idx  就具有这样的特点。它的值很容易预测：从 start 开始，然后伴随遍历而自增。

def enumerate_(iterable, start=0):  
    idx = start            # <- Initialisation.  
    for elem in iterable:    
        yield idx, elem      
        idx += 1           # <- (Predictable) Increment.

举一个我们之前看过的例子：

>>> list(enumerate_("rod"))
[(5, 'r'), (6, 'o'), (7, 'd')]

让我们把关注点放在由 idx 生成的数字上，如何生成单独生成它们呢？

>>> list(enumerate("rod", 5))
[(5, 'r'), (6, 'o'), (7, 'd')]

>>> ...
[5, 6, 7]

可以使用 list(range(...))：

>>> list(enumerate("rod", 5))
[(5, 'r'), (6, 'o'), (7, 'd')]

>>> list(range(5, 5 + 3))
[5, 6, 7]

内置的 range 函数可以实现 idx 一直在做的事：它从一个特定的初始值开始，递增生成连续的数字。

使用 range 对索引进行记录看起来不错，我们只需要确定传给 range的参数即可。起始值与 enumerate  的 start 值相同，而结束值应该是起始值加上可迭代对象的长度。

现在我们有了更好的方式来记录索引，只要同时遍历可迭代对象和 range ，就可以同时产生索引和元素：

def enumerate_(iterable, start=0):  
    idxs = range(start, start + len(iterable)) 
    for i in range(len(iterable)):     
        yield idxs[i], iterable[i]

不过，从一开始我们就已经使用过 for 循环，而且 range(len(...)) 好像并不适合出现在循环中。

目前，我们的 for 循环使用索引 i 可以同时遍历 idx 和可迭代对象的元素。在 python 中有一个内置的函数可以直接达成这个目的：即 zip 函数。通过使用它，可以实现同时遍历两个或多个可迭代对象，比如：

>>> list(zip(range(3), "rod"))
[(0, 'r'), (1, 'o'), (2, 'd')]

猜猜为什么 zip 也可以使用 list 调用？因为 zip 产生的也是生成器。

在知晓了 zip 函数可以达成索引和元素同时遍历的目标后，让我们更新一下刚刚的代码：

def enumerate_(iterable, start=0):  
    idxs = range(start, start + len(iterable))  
    for idx, elem in zip(idxs, iterable):    
        yield idx, elem

至此， 我们已经介绍了使用 for 循环配合 zip 函数实现两个可迭代对象的同步遍历。然而，这种看似正确的方式也可能导致执行错误。

可迭代对象，而非序列

enumerate_ 目前已经可以配合多种可迭代对象使用：

>>> list(enumerate_("rod"))
[(0, 'r'), (1, 'o'), (2, 'd')]

>>> list(enumerate_(["hello", "world", "!"]))
[(0, 'hello'), (1, 'world'), (2, '!')]

>>> list(enumerate_(range(0, 30, 10)))
[(0, 0), (1, 10), (2, 20)]

内置函数 enumerate  有一个很好的特性：只要输入的参数是一个可迭代对象， 它就会知道如何处理它。

让我们仔细观察，上述使用的参数都有一个长度属性：

• 字符串 "rod" 的长度为 3
• 列表 ["hello", "world", "!"] 的长度为 3
• range 对象 range(0, 30, 10) 的长度也为 3

这里重点不是它们的长度的值，而是对这些对象都可以使用内置的  len 函数：

>>> len("rod")
3

>>> len(["hello", "world", "!"])
3

>>> len(range(0, 30, 10))
3

在这一点上，我们的 enumerate_  函数功能并不完整。比如我们可以将  zip 传入 enumerate ：

>>> firsts = ["Harry", "Ron", "Hermione"]
>>> lasts = ["Potter", "Weasly", "Granger"]

>>> list(enumerate(zip(firsts, lasts)))
[(0, ('Harry', 'Potter')), (1, ('Ron', 'Weasly')), (2, ('Hermione', 'Granger'))]

但 enumerate_ 却不行：

>>> list(enumerate_(zip(firsts, lasts))) 
Traceback (most recent call last): 
  File "<stdin>", line 1, in <module> 
  File "<stdin>", line 2, in enumerate_
TypeError: object of type 'zip' has no len()

错误信息很明确地说明了问题：我们不能在 zip 对象上使用 len 函数。最重要的问题是，enumerate  应该适用于任何可迭代对象，但我们实现的 enumerate_ 仅对能计算长度的可迭代对象生效。

当遇到问题时，要思考这个问题是如何触发的。回顾刚才的代码可以得知：我们在计算 range 函数的结束值时用到了 len 函数。正如我们所见，一些可迭代对象我们并不知道它们的长度，因此不再能依靠 range 来确定遍历的起止区间。

我们需要实现一个新的索引生成方式，比如使用生成器：从一个给定的值开始，依次产生连续的数字。但是终止值该如何获取呢？

“并非所有美好的事情都需要结束”

为了解决刚遇到的问题，我们决定创建一个“永不停止”的生成器，它可能长这个样子：

def gen_indices(start):   
    idx = start  
    while True:   
        yield idx    
        idx += 1

完成函数定义之后，让我们试着使用它：

>>> for idx in gen_indices(0):
...     print(idx, end=" ")
... 
0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 
33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54
Traceback (most recent call last): 
  File "<stdin>", line 2, in <module>
KeyboardInterrupt

这里可通过 Ctrl + C 来中止程序运行。

看起来我们的无限生成器并不是很有用...不过这取决于我们如何使用它。

是时候介绍内置函数 next了。当我们使用生成器或一般的迭代器时，这是一个非常便利的函数。它的作用是：给定一个迭代器，然后返回该迭代器的下一项。

简单来说，我们可以用 next 获取生成器的下一项。比如定义一个从 42 开始计数的生成器：

>>> count_from_42 = gen_indices(42)
>>> count_from_42
<generator object gen_indices at 0x00000195DE3AA890>

现在我们想看一下 count_from_42 产生的元素。显然，使用 list 并不能生效，因为列表不能被定义为无限长度。因此，这里需要使用 next：

>>> count_from_42 = gen_indices(42)
>>> next(count_from_42)
42
>>> next(count_from_42)
43
>>> next(count_from_42)
44
>>> next(count_from_42
45

我们可以使用 next 手动管理索引的生成器，但实际上不需要这样做。zip 函数在使用时，伴随着参数迭代完毕将会自动停止，因此我们不必担心 gen_indices 是无限的。

def enumerate_(iterable, start=0):  
    for idx, elem in zip(gen_indices(start), iterable):   
        yield idx, elem

>>> list(enumerate_("rod", 5))
[(5, 'r'), (6, 'o'), (7, 'd')]

>>> list(enumerate_(zip(firsts, lasts)))
[(0, ('Harry', 'Potter')), (1, ('Ron', 'Weasly')), (2, ('Hermione', 'Granger'))]

值得再次强调的是：使用  gen_indices  而非 range ，效果会更好，因为前者的约束更少（换句话说，range 在使用时需要知道起始值和结束值，而  gen_indices  只需要知道起始值）。

再次记录索引

我们刚刚又引入了五行代码，显然，计划在第一次尝试时就达成目标的想法很天真。事实上，如果仔细观察这五行代码。会发现它与上一次记录索引的方式如出一辙：

def gen_indices(start):  
    idx = start    # <- set `idx`  
    while True:     
        yield idx   
        idx += 1   # <- always increment `idx`

如果你熟悉 C/C++、Java、JS 或其他语言，你应该会意识到 yield idx; idx += 1 这种写法是可以被替换的。比如使用 := 海象运算符^[1] ，可以实现在声明 idx 的同时对其实现增加一的操作：

idx := idx + 1

另外，由于 idx := idx + 1 是一个表达式，它可以和关键字 yield 配合使用，组成一个 yield 语句：

def gen_indices(start):   
    idx = start  
    while True:     
        yield (idx := idx + 1)

这样做以后，我们实现了将元素产生和自增结合到一条语句。除了索引值差 1 之外：

>>> list(enumerate_("rod"))
[(1, 'r'), (2, 'o'), (3, 'd')]

问题在于我们的赋值表达式 idx := idx + 1，它计算出了增量之后的值，但我们却忽略了设置 idx 的起始值。现在有两种方式来解决这个问题：

1、直接在 idx 声明时减一：

def gen_indices(start):   
    idx = start - 1  
    while True:     
        yield (idx := idx + 1)

不过这个看起来不是很优雅，因为对函数入参不够“尊重”；

2、在 yiled 时减一

def gen_indices(start):  
    idx = start  
    while True:    
        yield (idx := idx + 1) - 1

这种方式也不够优雅，它看起来就好像在撤销刚刚做的增量。

需要注意的是，编写更短的代码并不是 python 的目的。因此值得我们使用的，可能就是最初定义的那个版本：

def gen_indices(start):   
    idx = start  
    while True:     
        yield idx   
        idx += 1

如上文所述，这个版本会出现执行错误，那么正确的执行方式是怎样的？

寻找合适的工具

Python 是一门庞大的语言，尤其是它的标准库。所以每当我用复杂的方式实现简单逻辑时，我经常会选择去研究一下标准库是如何实现类似功能的。比如 itertools：

from itertools import count

itertools.count 就是一个合适的工具：

>>> help(count)

# ... 
 |  Equivalent to:
 |      def count(firstval=0, step=1): 
 |          x = firstval
 |          while 1:
 |              yield x 
 |              x += step

from itertools import count
def enumerate_(iterable, start=0):   
    for idx, elem in zip(count(start), iterable):  
        yield idx, elem

对于我们担心的场景，它依然正常：

>>> list(enumerate_("rod"))
[(0, 'r'), (1, 'o'), (2, 'd')]

你可能也认同，python 中总有一些新东西要学^[2]。但我的建议是多去了解一下标准库中的工具。通过正确使用它们，我们可以编写出更有表现力的代码。诚然，知道如何利用标准库，也是 python 程序员必备的技能。

从另一个可迭代对象中产生

现在我们不再被索引生成器牵制，可以仔细看看我们已经实现的 for 循环。它所产生的元素的准确性与我们期待的十分匹配。以至于我们好像不需要把 idx 和 elem 分开写，仅需要把它们放在一起即可：

>>> from itertools import count
>>> def enumerate_(iterable, start=0):
...     for idx_elem in zip(count(start), iterable):
...         yield idx_elem
...
>>> list(enumerate_("rod"))
[(0, 'r'), (1, 'o'), (2, 'd')]

也可以替换一下变量名：

from itertools import count
def enumerate_(iterable, start=0):  
    for t in zip(count(start), iterable):  
        yield t

我们的 for 循环中产生的元素都是从 zip 对象中取得的： enumerate_ 函数接收一个可迭代对象，使用参数构建另一个可迭代对象，最终，元素从这个可迭代对象中取出。这种模式非常常见，以至于有一个关键字，那就是 yield from：

>>> from itertools import count
>>> def enumerate_(it, start=0):
...     yield from zip(count(start), it)
...
>>> list(enumerate_("rod"))
[(0, 'r'), (1, 'o'), (2, 'd')]

这难道不是很酷吗？

挑战：<class 'enumerate'>

在结束本文之前，我们再注意一下另一个细节。直接调用内置的 enumerate， python 会将其视为一个类：

>>> enumerate
<class 'enumerate'>

而我们自己实现的 enumerate_ 是一个生成器函数：

>>> enumerate_
<function enumerate_ at 0x00000230C466EF70>

我们可以让 enumerate_ 模仿地更像一点吗？换句话说，我们可不可以将 enumerate_实现为一个类？再深入一点，我们是否可以将  enumerate_  实现为一个可生成惰性可迭代对象的类呢？显然是可以的，不过这将会是另一篇文章的主题，在这里不多赘述。

总结

在对内置的 enumerate 函数进行模拟实现的过程中，我们了解到：

• for 循环的习惯用法
• 可迭代的概念
• 什么是 生成器，以及如何使用 yield 定义 生成器
• zip 的目的
• 如何编写“无限”生成器
• 如何使用 next 手动从生成器中获取值
• 模块 itertools 及其 count 类
• 如何使用 yiled from 从另一个生成器中产生值

除了以上的技术知识以外，还需要强调的是，很多事情不是一蹴而就的。我们不害怕尝试，即使中途方向有些偏离，但至少学到了一些东西。我们不追求完美的解决方案，我们能做的只有不停的改进！

参考

• “Zip-up” Pydon't, https://mathspp.com/blog/pydonts/zip-up [last accessed 19-02-2022];
• “Bite-sized refactoring” Pydon't, https://mathspp.com/blog/pydonts/bite-sized-refactoring [last accessed 19-02-2022];
• “Why mastering Python is impossible, and why that's ok” Pydon't, https://mathspp.com/blog/pydonts/why-mastering-python-is-impossible [last accessed 19-02-2022];
• Twitter thread on enumerate, https://twitter.com/mathsppblog/status/1455444589603557378 [last accessed 19-02-2022];

参考资料