Fomalhaut
Fomalhaut

Reputation: 9785

How to repeat the body of a with-statement in Python?

I want to implement a way to repeat a section of code as many times as it's needed using a context manager only, because of its pretty syntax. Like this:

with try_until_success(attempts=10):
    command1()
    command2()
    command3()

The commands must be executed once if no errors happen. And they should be executed again if an error occurred, until 10 attempts has passed, if so the error must be raised. For example, it can be useful to reconnect to a data base. The syntax I represented is literal, I do not want to modify it (so do not suggest me to replace it with a kind of for of while statements).

Is there a way to implement try_until_success in Python to do what I want?

What I tried is:

from contextlib import contextmanager

@contextmanager
def try_until_success(attempts=None):
    counter = 0

    while True:
        try:
            yield
        except Exception as exc:
            pass
        else:
            break

        counter += 1
        if attempts is not None and counter >= attempts:
            raise exc

And this gives me the error:

RuntimeError: generator didn't stop after throw()

I know, there are many ways to reach what I need using a loop instead of with-statement or with the help of a decorator. But both have syntax disadvantages. For example, in case of a loop I have to insert try-except block, and in case of a decorator I have to define a new function.

I have already looked at the questions:

How do I make a contextmanager with a loop inside?

Conditionally skipping the body of Python With statement

They did not help in my question.

Upvotes: 7

Views: 1275

Answers (3)

Sam Mason
Sam Mason

Reputation: 16194

This goes against how context managers were designed to work, you'd likely have to resort to non-standard tricks like patching the bytecode to do this.

See the official docs on the with statement and the original PEP 343 for how they are expanded. It might help you understand why this isn't going to be officially supported, and maybe why other commenters are generally saying this is a bad thing to try and do.

As an example of something that might work, maybe try:

class try_until_success:
    def __init__(self, attempts):
        self.attempts = attempts
        self.attempt = 0
        self.done = False
        self.failures = []
    
    def __iter__(self):
        while not self.done and self.attempt < self.attempts:
            i = self.attempt
            yield self
            assert i != self.attempt, "attempt not attempted"
    
        if self.done:
            return
        
        if self.failures:
            raise Exception("failures occurred", self.failures)
        
    def __enter__(self):
        self.attempt += 1
    
    def __exit__(self, _ext, exc, _tb):
        if exc:
            self.failures.append(exc)
            return True
        
        self.done = True

for attempt in try_until_success(attempts=10):
    with attempt:
        command1()
        command2()
        command3()

you'd probably want to separate out the context manager from the iterator (to help prevent incorrect usage) but it sort of does something similar to what you were after

Upvotes: 5

Tomerikoo
Tomerikoo

Reputation: 19422

Is there a way to implement try_until_success in Python to do what I want?

Yes. You don't need to make it a context manager. Just make it a function accepting a function:

def try_until_success(command, attempts=1):
    for _ in range(attempts):
        try:
            return command()
        except Exception as exc:
            err = exc

    raise err

And then the syntax is still pretty clear, no for or while statements - not even with:

attempts = 10
try_until_success(command1, attempts)
try_until_success(command2, attempts)
try_until_success(command3, attempts)

Upvotes: 1

chepner
chepner

Reputation: 531718

The problem is that the body of the with statement does not run within the call to try_until_success. That function returns an object with a __enter__ method; that __enter__ method calls and returns, then the body of the with statement is executed. There is no provision for wrapping the body in any kind of loop that would allow it to be repeated once the end of the with statement is reached.

Upvotes: 2

Related Questions