Как я могу убедиться, что генератор правильно закрыт?

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 метод. Но его подпись не обещает этого. Что, если его реализация проще, чем я дал ранее?

>>> 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 управлять вашими ресурсами, не беспокоясь об обработке релиза. Если вы чувствуете себя смелым, вы можете даже создать свой собственный менеджер контекста класс с __enter__ и __exit__ функцией для управления вашим ресурсом.

Я думаю, что ключевым моментом здесь является то, чтотак как ваш генератор должен управлять ресурсом, вы должны либо использовать оператор with, либо всегда закрывать его после него, так же, как f = open(...) всегда должен следовать с f.close()

Answer 2

Один из вариантов - использовать Generator тип , чтобы правильно сигнализировать, что вы возвращаете Generator.