Decorator which re-runs a function if it exited with a specified code

exit_retry.py

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from sys import _ExitCode

def exit_retry(tries: int, code: "_ExitCode | None" = None, pause: float = 0):
    """Retry function call after sys.exit().

    :param retries: number of retries
    :type retries: int
    :param code: exit code to trigger retry (None=any)
    :type code: str | int | None
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_code: _ExitCode = None
            for i in range(tries):
                try:
                    if i:  # pause between iterations
                        print(f"Retry after {pause}s...")
                        sleep(pause)
                    return func(*args, **kwargs)
                except SystemExit as e:
                    last_code = e.code
                    if not code or e.code == code:
                        continue
                    raise
            sys.exit(last_code)
        return wrapper
    return decorator