Пользовательский итератор или генератор для выполнения распродажи после перерыва

Question

Я реализую настроенное поведение цикла, где мне нужно, чтобы что-то происходило при входе в цикл, при каждом запуске цикла, на каждом конце цикла и при выходе из области цикла. Пока что это очень просто в Python (2.7):

def my_for(loop_iterable):
    enter_loop()
    for i in loop_iterable:
        loop_start()
        yield i
        loop_end()
    exit_loop()

for i in my_for([1, 2, 3]):
    print "i: ", i
    if i == 2:
        break

Проблема у меня в получении loop_end() а также exit_loop() выполнить после break, Я решил это, определив другую функцию, которую пользователь должен поставить перед перерывом:

def break_loop():
    loop_end()
    exit_loop()

for i in my_for([1, 2, 3]):
    print "i: ", i
    if i == 2:
        break_loop()
        break

Но мне бы очень хотелось, чтобы пользователь не забыл добавить эту строку. Я думаю, что если я переписываю функцию генератора как класс итератора, возможно, есть способ выполнить код на break?

Между прочим, continue работает просто отлично как есть!

Answer 1

Вы можете использовать контекстный менеджер:

class Looper(object):
    def __init__(self, iterable):
        self.iterable = iterable
        self.need_to_end = False

    def __enter__(self):
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.exit_loop()
        # Handle exceptions or swallow them by returning True

    def enter_loop(self):
        print 'enter_loop'

    def loop_start(self):
        self.need_to_end = True
        print 'loop_start'

    def loop_end(self):
        self.need_to_end = False
        print 'loop_end'

    def exit_loop(self):
        if self.need_to_end:
            self.loop_end()

        print 'exit_loop'

    def __iter__(self):
        self.enter_loop()

        for i in self.iterable:
            self.loop_start()
            yield i
            self.loop_end()

Ваш код получился бы немного длиннее, но вы можете разобраться с исключениями и другими вещами более чисто:

with Looper([1, 2, 3, 4, 5, 6]) as loop:
    for i in loop:
        print i

        if i == 2:
            continue
        elif i == 3:
            break

Это работает так, как вы ожидаете:

enter_loop
loop_start
1
loop_end
loop_start
2
loop_end
loop_start
3
loop_end
exit_loop

Answer 2

Вы можете использовать __enter__ а также __exit__ магическая функция, определяя ее в классе. Чтобы сделать звонок, вы можете использовать его с with, __enter__ метод будет вызван перед выполнением кода в with и когда блок with блок выйдет, __exit__ функция будет вызвана. Например:

>>> class MyTestWrapper(object):
...     def __enter__(self):
...         print 'I am in __enter__'
...     def __exit__(self, type, value, traceback):
...         print 'I am in __exit__'
...
>>> with MyTestWrapper() as s:
...     print 'My Loop Logic'
...
I am in __enter__
My Loop Logic
I am in __exit__

Теперь, чтобы сделать его итератором, нужно определить __iter__ функция. При этом вы можете назвать это как итератор. Обновленный код будет:

>>> class MyIterator(object):
...     def __init__(self, iterable):
...         self.iterable = iterable
...         self.need_to_end = False
...     def __enter__(self):
...         print 'I am in __enter__'
...         return self
...     def __exit__(self, type, value, traceback):
...         self.loop_end()
...         print 'I am in __exit__'
...     def loop_start(self):
...         self.need_to_end = True
...         print 'Starting Loop . . . '
...     def loop_end(self):
...         self.need_to_end = False
...         print 'Ending Loop . . . '
...     def __iter__(self):
...         for i in self.iterable:
...             self.loop_start()
...             yield i
...             self.loop_end()
...
>>> with MyIterator([1,2,3, 4]) as my_iterator:
...     for i in my_iterator:
...         print 'I am: ', i
...         if i == 2:
...             break
...
I am in __enter__
Starting Loop . . .
I am:  1
Ending Loop . . .
Starting Loop . . .
I am:  2
Ending Loop . . .
I am in __exit__

Обратитесь к Менеджеру контекста Python для получения дополнительной информации об этих и других встроенных функциях.