python - 如何确保发电机正确关闭？

Question

考虑具有以下签名的库函数:

from typing import Iterator

def get_numbers() -> Iterator[int]:
    ...

让我们看一些使用它的简单代码:

for i in get_numbers():
    print(i)

到目前为止没有什么有趣的事情。但假设我们不关心偶数。仅限奇数，例如我们:

for i in get_numbers():
    if i & 1 == 0:
        raise ValueError("Ew, an even number!")
    print(i)

现在让我们尝试一下 get_numbers 的实现:

def get_numbers() -> Iterator[int]:
    yield 1
    yield 2
    yield 3

这里没什么有趣的。运行我们的小 for 的结果几乎符合我们的预期:

>>> for i in get_numbers():
  2     if i & 1 == 0:
  3         raise ValueError("Ew, an even number!")
  4     print(i)
1
Traceback (most recent call last):
  File "<stdin>", line 3, in <module>
ValueError: Ew, an even number!

Ew, an even number!
>>>

如果 get_numbers 有更简单的实现，我们会得到完全相同的结果:

def get_numbers() -> Iterator[int]:
    return iter([1, 2, 3])

但是我们假设 get_numbers 需要保留为生成器，因为它管理一些资源。

def get_numbers() -> Iterator[int]:
    acquire_some_resource()
    try:
        yield 1
        yield 2
        yield 3
    finally:
        release_some_resource()

就我们的目的而言，我们将管理的资源只是打印在屏幕上的文本:

def acquire_some_resource() -> None:
    print("generating some numbers")

def release_some_resource() -> None:
    print("done generating numbers")

我们的输出仍然是可预测的:

>>> for i in get_numbers():
  2     if i & 1 == 0:
  3         raise ValueError("Ew, an even number!")
  4     print(i)
generating some numbers
1
done generating numbers
Traceback (most recent call last):
  File "<stdin>", line 3, in <module>
ValueError: Ew, an even number!

Ew, an even number!
>>>

但是如果我们不能使用简单的 for 循环怎么办？例如，如果我们想忽略第一个数字怎么办？ (假设 itertools.islice 不是一个东西。)

>>> it = get_numbers()
  2 next(it, None)
  3 for i in it:
  4     if i & 1 == 0:
  5         raise ValueError("Ew, an even number!")
  6     print(i)
generating some numbers
Traceback (most recent call last):
  File "<stdin>", line 5, in <module>
ValueError: Ew, an even number!

Ew, an even number!
>>>

注意到什么了吗？我们获得了我们的资源，如“生成一些数字”文本所示，但我们从未发布过它。

正确的做法是确保发电机关闭:

>>> it = get_numbers()
  2 try:
  3     next(it, None)
  4     for i in it:
  5         if i & 1 == 0:
  6             raise ValueError("Ew, an even number!")
  7         print(i)
  8 finally:
  9     it.close()
generating some numbers
done generating numbers
Traceback (most recent call last):
  File "<stdin>", line 6, in <module>
ValueError: Ew, an even number!

Ew, an even number!
>>>

这种方法的问题在于，它假设 get_numbers() 返回一个生成器，因此有一个 close 方法。但它的签名并没有 promise 这一点。如果它的实现比我之前给出的更简单怎么办？

>>> def get_numbers() -> Iterator[int]:
  2     return iter([1, 2, 3])
  3 
  4 it = get_numbers()
  5 try:
  6     next(it, None)
  7     for i in it:
  8         if i & 1 == 0:
  9             raise ValueError("Ew, an even number!")
 10         print(i)
 11 finally:
 12     it.close()
Traceback (most recent call last):
  File "<stdin>", line 12, in <module>
AttributeError: 'list_iterator' object has no attribute 'close'

'list_iterator' object has no attribute 'close'
>>>

因此，这里要做的正确事情是相当乏味的:

it = get_numbers()
try:
    next(it, None)
    for i in it: 
        if i & 1 == 0: 
            raise ValueError("Ew, an even number!") 
        print(i) 
finally: 
    if hasattr(it, "close"): 
        it.close()

我可以将其包装在上下文管理器中以使其更简单，但感觉我正在做语言应该为我做的事情，或者至少，被调用者应该关心自己，而不是调用者.

有没有更简单的方法来处理这个问题？

Answer 1

正如我的评论所提到的，正确构建此结构的一种方法是使用 contextlib.contextmanager装饰你的发电机:

from typing import Iterator
import contextlib

@contextlib.contextmanager
def get_numbers() -> Iterator[int]:
    acquire_some_resource()
    try:
        yield iter([1, 2, 3])
    finally:
        release_some_resource()

然后当你使用生成器时:

with get_numbers() as et:
    for i in et:
        if i % 2 == 0:
            raise ValueError()
        else:
            print(i)

结果:

generating some numbers
1
done generating numbers
Traceback (most recent call last):
  File "<pyshell#64>", line 4, in <module>
    raise ValueError()
ValueError

这允许 contextmanager 装饰器为您管理资源，而无需担心处理发布。如果您有勇气，您甚至可以build your own context manager具有 __enter__ 和 __exit__ 函数的类来处理您的资源。

我认为这里的关键要点是，由于您的生成器需要管理资源，因此您应该使用 with 语句，或者始终在之后关闭它，就像 f = open(...) 后面应始终跟有 f.close()