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

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实现更简单，我们将获得完全相同的结果：

>>> 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])

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

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

我们的输出仍然可以预测：

>>> 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方法。但是它的签名并不能保证这一点。如果它的实现是我之前给出的更简单的实现，该怎么办？

>>> 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'
>>>

因此，在这里要做的正确的事情很繁琐：

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

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

我可以将其包装在上下文管理器中以使其更简单，但是感觉就像我在做某种语言应该为我做的事情，或者至少是被调用方应该与自身而不是与调用方有关。

有没有更简单的方法来解决这个问题？

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装饰器为您管理资源，而不必担心处理发布。如果您有勇气，您甚至可以构建自己的上下文管理器类，__enter__并使用__exit__它的函数来处理您的资源。

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