Python Retry manager with a generator and context manager

retry_context.py

class RetryContext:
    """Retry context manager creator

    created by TheRolf"""

    def __init__(self, retries=1, exceptions=(Exception,), delay=0.0, on_fail=None) -> None:
        """Creates retry context creator

        Args:
            retries (int, optional): Number of retries allowed. Defaults to 1.
            exceptions (tuple|Exception, optional): Exception class or exception class tuple to catch. Defaults to all exceptions (Exception,).
            delay (float, optional): Delay in seconds between failed runs. Defaults to 0.
            on_fail (function, optional): Function to execute on fail, takes the error as argument.
        """
        self.retries = retries
        self.delay = delay
        self.exceptions = exceptions
        self.on_fail = on_fail

        # value used to capture error and detect run failure
        self._captured_error = None
        self.attempt_nb = 0

    @property
    def captured_error(self):
        """Returns the last captured error

        Returns:
            (Exception|None): The captured error or None if no error occured
        """
        return self._captured_error

    @captured_error.setter
    def captured_error(self, value):
        assert value is not None
        # increases attempt number and set captured error in object
        self._captured_error = value
        self.attempt_nb += 1

    def launch(self):
        """Starts a generator of retry runs context managers

        Yields:
            RetryContext.RetryRun: Retry run context manager class
        """
        # To make object and method reusable, clear
        self._captured_error = None
        self.attempt_nb = 0

        errored = True # exit condition
        while errored and self.attempt_nb <= self.retries:
            # give Context manager for run
            yield self.RetryRun(self)

            # update exit condition
            errored = self._captured_error is not None

            #! clear last captured error
            self._captured_error = None

            # if failed, do the failure action
            if errored:
                if self.on_fail:
                    self.on_fail(self.attempt_nb)
                # no need for a delay if all attempts failed
                if self.delay > 0 and self.attempt_nb <= self.retries:
                    import time
                    time.sleep(self.delay)


    class RetryRun:
        """Retry run context manager class"""

        def __init__(self, ctx) -> None:
            """
            Args:
                ctx (RetryContext): parent context builder instance
            """
            self._ctx = ctx

        def __enter__(self):
            # remove any type of reference by multiplying by one
            # give in the context the attempt number because why not
            return int(self._ctx.attempt_nb * 1)

        def __exit__(self, err_type, err_value, _err_traceback):
            # no error
            if err_type is None:
                return False  # propagate nothing, no error happened, exit normally

            # Exception type not in the specified exceptions, exit normally
            if not issubclass(err_type, self._ctx.exceptions):
                return False # propagate the unhandled exception

            # update the captured error and attempt number
            self._ctx.captured_error = err_value

            # attempt number can be equal to retries, next run will be last attempt
            if self._ctx.attempt_nb <= self._ctx.retries:
                return True  # suppress the exception

            return False  # propagate exception

retry_context_example.py

print("First retry context launch: passing on second run")
retry_context = RetryContext(retries=1, exceptions=(ValueError,))
for retry_run in retry_context.launch(): # generator
    with retry_run as run_nb: # context manager
        print(f"Run {run_nb}")
        if run_nb == 0:
            print("Failed on first run")
            raise ValueError("Oh no")
        else:
            print("Passed on next run")
print("Second retry passed successfully")

print("Second retry context launch: failing after second run at ValueError")
try:
    for retry_run in retry_context.launch(): # generator
        with retry_run as run_nb: # context manager
            print(f"Run {run_nb}")
            raise ValueError(f"Failed on attempt {run_nb}")
except ValueError as e:
    assert isinstance(e, ValueError), "Class shall match"
    assert str(e) == "Failed on attempt 1", "Error message shall match"
    print("Second retry failed successfully")

retry_context_machine_example.py

class Machine:
    broken: bool
    def __init__(self) -> None:
        self.broken = False
        print("MACHINE ONLINE")
    def explode(self):
        print("KABOOM")
        self.broken = True
    def run(self):
        if self.broken:
            raise RuntimeError("Machine is broken")
        print("WORK WORK")
    def repair(self):
        print("BIP BOOP")
        self.broken = False

machine = Machine()
machine.explode()

for retry_run in RetryContext(exceptions=RuntimeError, on_fail=lambda err: machine.repair()).launch(): # generator
    with retry_run: # context manager
        machine.run()