Adam S. Adam S. - 3 months ago 14
Python Question

custom iterator or generator to execute closeout after break

I'm implementing customized loop behavior, where I need things to happen on entering the loop, at every loop start, at every loop end, and on exiting the loop area. So far this is beautifully simple in 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


The problem I'm having is in getting
loop_end()
and
exit_loop()
to execute after the
break
. I have solved this by defining another function, which the user must put before 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


But I would really like not to have the user have to remember to add that line. I think if I re-write the generator function as an iterator class, maybe there is a way to still execute code on a
break
?

Incidentally,
continue
works just fine as is!

Answer

You could use a 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()

Your code would get a little longer, but you can deal with exceptions and other things more cleanly:

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

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

It works as you'd expect:

enter_loop
loop_start
1
loop_end
loop_start
2
loop_end
loop_start
3
loop_end
exit_loop