使用生成器函数自定义迭代

Beginner

This tutorial is from open-source community. Access the source code

简介

本节将探讨如何使用生成器函数来自定义迭代。

一个问题

假设你想要创建自己的自定义迭代模式。

例如,一个倒计时。

>>> for x in countdown(10):
...   print(x, end=' ')
...
10 9 8 7 6 5 4 3 2 1
>>>

有一个简单的方法可以做到这一点。

生成器

生成器是一种定义迭代的函数。

def countdown(n):
    while n > 0:
        yield n
        n -= 1

例如:

>>> for x in countdown(10):
...   print(x, end=' ')
...
10 9 8 7 6 5 4 3 2 1
>>>

生成器是任何使用 yield 语句的函数。

生成器的行为与普通函数不同。调用生成器函数会创建一个生成器对象。它不会立即执行函数。

def countdown(n):
    ## 添加了一条打印语句
    print('Counting down from', n)
    while n > 0:
        yield n
        n -= 1
>>> x = countdown(10)
## 没有打印语句
>>> x
## x 是一个生成器对象
<generator object at 0x58490>
>>>

函数仅在调用 __next__() 时执行。

>>> x = countdown(10)
>>> x
<generator object at 0x58490>
>>> x.__next__()
Counting down from 10
10
>>>

yield 产生一个值,但会暂停函数执行。函数会在下次调用 __next__() 时恢复执行。

>>> x.__next__()
9
>>> x.__next__()
8

当生成器最终返回时,迭代会引发一个错误。

>>> x.__next__()
1
>>> x.__next__()
Traceback (most recent call last):
File "<stdin>", line 1, in? StopIteration
>>>

观察:生成器函数实现了与 for 语句在列表、元组、字典、文件等上使用的相同低级协议。

练习 6.4:一个简单的生成器

如果你发现自己想要自定义迭代,那么你应该始终考虑生成器函数。它们很容易编写——创建一个执行所需迭代逻辑的函数,并使用 yield 来发出值。

例如,试试这个在文件中搜索包含匹配子字符串的行的生成器:

>>> def filematch(filename, substr):
        with open(filename, 'r') as f:
            for line in f:
                if substr in line:
                    yield line

>>> for line in open('portfolio.csv'):
        print(line, end='')

name,shares,price
"AA",100,32.20
"IBM",50,91.10
"CAT",150,83.44
"MSFT",200,51.23
"GE",95,40.37
"MSFT",50,65.10
"IBM",100,70.44
>>> for line in filematch('portfolio.csv', 'IBM'):
        print(line, end='')

"IBM",50,91.10
"IBM",100,70.44
>>>

这有点意思——你可以将一堆自定义处理隐藏在一个函数中,并使用它来为 for 循环提供数据。下一个示例将看看一个更不寻常的情况。

练习 6.5:监控流数据源

生成器可以是一种监控实时数据源(如日志文件或股票市场数据馈送)的有趣方式。在本部分中,我们将探讨这个想法。首先,请仔细按照以下说明操作。

程序 stocksim.py 是一个模拟股票市场数据的程序。作为输出,该程序会不断将实时数据写入文件 stocklog.csv。在一个单独的命令窗口中,进入该目录并运行此程序:

$ python3 stocksim.py

如果你使用的是 Windows,只需找到 stocksim.py 程序并双击运行它。现在,先不管这个程序(让它运行就行)。在另一个窗口中,查看模拟器正在写入的文件 stocklog.csv。你应该会看到每隔几秒就会有新的文本行添加到文件中。同样,就让这个程序在后台运行——它会运行几个小时(你无需担心它)。

一旦上述程序运行起来,让我们编写一个小程序来打开该文件,定位到文件末尾,并监视新的输出。创建一个文件 follow.py 并将以下代码放入其中:

## follow.py
import os
import time

f = open('stocklog.csv')
f.seek(0, os.SEEK_END)   ## 将文件指针从文件末尾移动 0 个字节

while True:
    line = f.readline()
    if line == '':
        time.sleep(0.1)   ## 短暂休眠并重试
        continue
    fields = line.split(',')
    name = fields[0].strip('"')
    price = float(fields[1])
    change = float(fields[4])
    if change < 0:
        print(f'{name:>10s} {price:>10.2f} {change:>10.2f}')

如果你运行该程序,你将看到一个实时股票报价器。在底层,这段代码有点类似于用于监视日志文件的 Unix tail -f 命令。

注意:在这个示例中使用 readline() 方法有点不寻常,因为它不是从文件中读取行的常用方式(通常你会使用 for 循环)。然而,在这种情况下,我们使用它来反复探测文件末尾,看看是否有更多数据被添加(readline() 要么返回新数据,要么返回一个空字符串)。

练习 6.6:使用生成器生成数据

如果你查看练习 6.5 中的代码,代码的第一部分在生成数据行,而 while 循环末尾的语句在消费数据。生成器函数的一个主要特性是你可以将所有数据生成代码移动到一个可复用的函数中。

修改练习 6.5 中的代码,使得文件读取由生成器函数 follow(filename) 来执行。确保以下代码能够正常工作:

>>> for line in follow('stocklog.csv'):
          print(line, end='')

... 这里应该会看到输出的行...

将股票报价器代码修改成如下形式:

if __name__ == '__main__':
    for line in follow('stocklog.csv'):
        fields = line.split(',')
        name = fields[0].strip('"')
        price = float(fields[1])
        change = float(fields[4])
        if change < 0:
            print(f'{name:>10s} {price:>10.2f} {change:>10.2f}')

练习 6.7:关注你的投资组合

修改 follow.py 程序,使其监视股票数据流,并仅打印出投资组合中那些股票的信息的报价器。例如:

if __name__ == '__main__':
    import report

    portfolio = report.read_portfolio('portfolio.csv')

    for line in follow('stocklog.csv'):
        fields = line.split(',')
        name = fields[0].strip('"')
        price = float(fields[1])
        change = float(fields[4])
        if name in portfolio:
            print(f'{name:>10s} {price:>10.2f} {change:>10.2f}')

注意:要使其正常工作,你的 Portfolio 类必须支持 in 运算符。请参阅练习 6.3,并确保你实现了 __contains__() 运算符。

讨论

这里刚刚发生了一些非常强大的事情。你将一种有趣的迭代模式(在文件末尾读取行)转移到了它自己的一个小函数中。follow() 函数现在成为了一个完全通用的实用工具,你可以在任何程序中使用它。例如,你可以用它来监视服务器日志、调试日志以及其他类似的数据源。这相当酷。

总结

恭喜你!你已经完成了“自定义迭代”实验。你可以在 LabEx 中练习更多实验来提升你的技能。