Python3: Telling a context manager to execute the code block more than once?-CodePudding

I'm wondering whether something is possible under python3. I would like to write a context manager which would execute its code block more than once. This doesn't seem to fit into the standard python3 context manger semantics, but I'm wondering whether I'm missing something. I want something like the following:

CtxMgr(repeat_count=N, ... other args ...):
    # This is the code block.
    # This code block does whatever it needs to do,
    # and this context manager makes sure that it
    # runs `repeat_count` times.

I'm guessing that if such a thing is possible, it might be implementable somehow under a structure similar to this:

class CtxMgr(contextlib.AbstractContextManager):
    def __init__(self, repeat_count=1, *args, **kwargs):
        self.repeat_count = repeat_count
        ## etc.
    def __enter__(self):
        ## etc.
    def __exit__(self, *args):
        ## etc.
    ## What other method(s) would need to be implemented here?

I know that I can do the following, but I don't want to do it this way. I'm looking for a way to do this solely as a context manager.

# What I don't want to do ...

def func(... ...):
    # whatever

CtxMgr(repeat_count=3) as f:
    for _ in range(0, f.repeat_count):
        func(... ...)

In other words, I'm wondering whether this could somehow be accomplished in python3 without the auxiliary function func, and for whatever code block follows the context manager instantiation, that code block would be automatically executed repeat_count times.

Thank you in advance.

CodePudding user response：

As I mentioned in my comment, this is not currently possible with a context manager. However, you can use a decorator:

import functools
import time
from typing import Collection, Type, Union


def retry_n_times(n: int, *, sleep_secs: float = 0.0, retry_on: Union[Type[Exception], Collection[Type[Exception]]] = Exception):
    """
    Retry the decorated function up to `n` times.  Raise a RuntimeError if all attempts fail.
    If `sleep_secs` is specified, sleep that long after each attempt.
    If `retry_on` is specified, only retry on these exception types - other exceptions will escape as normal.
      By default, retry on any exception.
    """
    if n <= 0:
        raise ValueError("'n' must be >= 1")

    if not isinstance(retry_on, Collection):
        retry_on = [retry_on]
    retry_on = tuple(retry_on)

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for _ in range(n):
                try:
                    return func(*args, **kwargs)
                except retry_on as exc:
                    last_exception = exc  # prevents "local variable 'exc' referenced before assignment"
                    if sleep_secs:
                        time.sleep(sleep_secs)
            else:
                raise RuntimeError(f'{func.__name__} failed {n} times') from last_exception

        return wrapper
    return decorator

Example usage:

count = 0

@retry_n_times(3)
def succeed_on_second_try():
    count  = 1
    if count == 1:
        raise ValueError('Intentional failure')

succeed_on_second_try()
assert count == 2