自定义迭代器或生成器以在中断后执行收尾

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 之后执行。我通过定义另一个函数解决了这个问题，用户必须在 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

但我真的希望用户不必记住添加该行。我认为如果我将生成器函数重写为迭代器 class，也许仍然可以在 break?

上执行代码

顺便说一下，continue 照原样工作得很好！

Answer 1

您可以通过在class中定义来使用__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's Context Manager。

Answer 2

您可以使用 context manager:

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

自定义迭代器或生成器以在中断后执行收尾

custom iterator or generator to execute closeout after break

python

iterator

loops

generator

break