Как мне сделать контекстный менеджер с циклом внутри?

Я хочу что-то вроде этого:

from contextlib import contextmanager

@contextmanager
def loop(seq):
    for i in seq:
        try:
            do_setup(i)
            yield # with body executes here
            do_cleanup(i)
        except CustomError as e:
            print(e)

with loop([1,2,3]):
    do_something_else()
    do_whatever()

Но contextmanager не работает, потому что ожидает, что генератор выдаст ровно один раз.

Причина, по которой я хочу этого, заключается в том, что я в основном хочу создать свой собственный цикл for. У меня есть модифицированный IPython, который используется для управления тестовым оборудованием. Очевидно, что это полноценный Python REPL, но большую часть времени пользователь просто вызывает предопределенные функции (аналогично приглашению Bash), и от пользователя не ожидается, что он программист или знаком с Python. Должен быть способ зацикливания некоторого произвольного кода с настройкой/очисткой и обработкой исключений для каждой итерации, и он должен быть примерно таким же простым, как приведенный выше оператор with.


person jpkotta    schedule 17.04.2015    source источник


Ответы (2)


Я думаю, что генератор работает лучше здесь:

def loop(seq):
    for i in seq:
        try:
            print('before')
            yield i  # with body executes here
            print('after')
        except CustomError as e:
            print(e)

for i in loop([1,2,3]):
    print(i)
    print('code')

дам:

before
1
code
after
before
2
code
after
before
3
code
after

Python входит и выходит из блока with только один раз, поэтому у вас не может быть логики в шагах входа/выхода, которые будут выполняться неоднократно.

person Simeon Visser    schedule 17.04.2015
comment
Старый вопрос, но я не понимаю, как генератор эквивалентен диспетчеру контекста. ihmo, одно из основных преимуществ кода OP заключается в том, что try...except перехватывает исключения в коде, которому уступили (например, do_whatever()).. это решение не делает. Любые идеи о том, как это получить? Типичным сценарием будет некоторая логика повторных попыток - person ttyridal; 05.02.2019
comment
@ttyridal: это зависит от того, где вы вызываете исключение. Если это происходит в генераторе, то все в порядке. Если это произойдет for i in loop... снаружи, оно не будет поймано, и вам нужно поймать его там (возможно, с помощью менеджера контекста). Кроме того, contextlib.contextmanager создает специальный генератор, который может возвращать только один раз, а мне нужна была штука, которая бы возвращала значение для каждого элемента последовательности. Вам нужен как менеджер контекста, так и генератор, чтобы быть полным решением (я опубликовал это как еще один ответ). - person jpkotta; 12.03.2019

Более полный ответ, если исключение может произойти вне генератора:

from contextlib import contextmanager

class CustomError(RuntimeError):
    pass

@contextmanager
def handle_custom_error():
    try:
        yield
    except CustomError as e:
        print(f"handled: {e}")

def loop(seq):
    for i in seq:
        try:
            print('before')
            if i == 0:
                raise CustomError("inside generator")
            yield i # for body executes here
            print('after')
        except CustomError as e:
            print(f"handled: {e}")

@handle_custom_error()
def do_stuff(i):
    if i == 1:
        raise CustomError("inside do_stuff")
    print(f"i = {i}")

for i in loop(range(3)):
    do_stuff(i)

Выход:

before
handled: inside generator
before
handled: inside do_stuff
after
before
i = 2
after
person jpkotta    schedule 11.03.2019
comment
Это отличное решение. Однако для этого необходимо, чтобы код do_stuff находился в отдельном декорированном методе. Как насчет решения, которое позволяет использовать более простой в использовании подход, подобный этому: for i in loop(range(3)): if i == 1: raise CustomError("inside code") Есть ли такой способ справиться с этим в python без необходимости использовать оба for с вложенным with в пользовательском коде? (Я думаю, что блоки кода ruby ​​​​допустили бы такое, отсюда и мой вопрос) - person Mirek; 04.04.2019
comment
По сути, я пытаюсь улучшить свою тестовую среду, чтобы разработчики QA могли легко использовать одну простую структуру для повторного выполнения кода (а не только метода), который может вызвать AssertionError (или какой-либо другой пользовательский кортеж исключений, поэтому генератор и менеджер контекста должны совместно использовать некоторая область, чтобы знать, какое исключение обрабатывать). - person Mirek; 04.04.2019
comment
@Mirek Я не знаю, есть ли способ сделать это. На самом деле это то, что я хотел сделать в первую очередь, но мой ответ достаточно хорош для моих целей (все ошибки, которые я хочу поймать, находятся в предопределенных функциях, и они не могут происходить вне этих функций). Чтобы сделать то, что вы хотите, я мог бы начать с рассмотрения того, как работает fuckit.py. - person jpkotta; 05.04.2019