Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How do I make a contextmanager with a loop inside?

I want something like this:

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()

But contextmanager doesn't work because it expects the generator to yield exactly once.

The reason why I want this is because I basically want to make my own custom for loop. I have a modified IPython that is used to control test equipment. It's obviously a full Python REPL, but most of the time the user is just calling predefined functions (similar to Bash prompt), and the user is not expected to be a programmer or familiar with Python. There needs to be a way to loop over some arbitrary code with setup/cleanup and exception handling for each iteration, and it should be about as simple to type as the above with statement.

like image 941
jpkotta Avatar asked Apr 17 '15 20:04

jpkotta


People also ask

What is ExitStack?

An ExitStack is (as the name suggests) a stack of clean-up functions. Adding a callback to the stack is the equivalent of calling Go's defer statement. However, clean-up functions are not executed when the function returns, but when execution leaves the with block - and until then, the stack can also be emptied again.

What is Contextlib?

Overview. The contextlib module of Python's standard library provides utilities for resource allocation to the with statement. The with statement in Python is used for resource management and exception handling. Therefore, it serves as a good Context Manager.

What does __ enter __ do in Python?

__enter__() is provided which returns self while object. __exit__() is an abstract method which by default returns None . See also the definition of Context Manager Types. New in version 3.6.


2 Answers

A more complete answer, for if the exception might happen outside the generator:

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)

Output:

before
handled: inside generator
before
handled: inside do_stuff
after
before
i = 2
after
like image 188
jpkotta Avatar answered Sep 20 '22 17:09

jpkotta


I think a generator works better here:

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')

will give:

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

Python enters and exits a with block only once so you can't have logic int the enter / exit steps that would be done repeatedly.

like image 40
Simeon Visser Avatar answered Sep 16 '22 17:09

Simeon Visser