Add retry utility for handling function retries (#278)

* add retry function

* make timeout optional

* review docstring

* add pytest

* remove logging_callback and add optional keyword arguments for on_failed

* allow for return values

* restructure retry

* break pytest into multiple short tests, use `pytest.raises`

---------

Co-authored-by: Falko Schindler <[email protected]>

Author: pascalzauberzeug

rosys/run.py

@@ -7,9 +7,11 @@
 import uuid
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
+from dataclasses import dataclass
 from functools import wraps
+from inspect import signature
 from pathlib import Path
-from typing import ParamSpec, TypeVar
+from typing import Any, ParamSpec, TypeVar
 
 from nicegui import run
 
@@ -143,3 +145,37 @@ def tear_down() -> None:
         _kill(process)
     running_sh_processes.clear()
     log.info('teardown complete.')
+
+
+@dataclass(slots=True, kw_only=True, frozen=True)
+class OnFailedArguments:
+    attempt: int
+    max_attempts: int
+
+
+async def retry(func: Callable, *,
+                max_attempts: int = 3,
+                max_timeout: float | None = None,
+                on_failed: Callable | None = None) -> Any:
+    """Call a function repeatedly until it succeeds or reaches the maximum number of attempts.
+
+    :param func: A function to retry
+    :param max_attempts: Maximum number of attempts
+    :param max_timeout: Optional maximum time in seconds to wait per attempt
+    :param on_failed: Optional callback to execute after each failed attempt (optional argument of type ``OnFailedArguments``)
+    :return: Result of the called function
+    :raises RuntimeError: If all attempts fail
+    """
+    for attempt in range(max_attempts):
+        try:
+            return await asyncio.wait_for(func(), timeout=max_timeout)
+        except Exception:
+            if on_failed is None:
+                continue
+            if signature(on_failed).parameters:
+                result = on_failed(OnFailedArguments(attempt=attempt, max_attempts=max_attempts))
+            else:
+                result = on_failed()
+            if asyncio.iscoroutinefunction(on_failed):
+                await result
+    raise RuntimeError(f'Running {func.__name__} failed.')

tests/test_run.py

@@ -0,0 +1,53 @@
+import pytest
+
+from rosys import run
+
+
+@pytest.mark.asyncio
+async def test_retry():
+    async def func() -> str:
+        events.append('call')
+        return 'success'
+
+    events: list[str] = []
+    result = await run.retry(func)
+    assert result == 'success'
+    assert events == ['call']
+
+
+@pytest.mark.asyncio
+async def test_retry_failed():
+    async def func() -> None:
+        events.append('call')
+        raise ValueError()
+
+    events: list[str] = []
+    with pytest.raises(RuntimeError):
+        await run.retry(func)
+    assert events == ['call'] * 3
+
+
+@pytest.mark.asyncio
+async def test_retry_failed_with_on_failed():
+    async def func() -> None:
+        events.append('call')
+        raise ValueError()
+
+    events: list[str] = []
+    with pytest.raises(RuntimeError):
+        await run.retry(func, on_failed=lambda: events.append('failed'))
+    assert events == ['call', 'failed'] * 3
+
+
+@pytest.mark.asyncio
+async def test_retry_failed_with_on_failed_and_args():
+    async def func() -> None:
+        events.append('call')
+        raise ValueError()
+
+    events: list[str] = []
+    with pytest.raises(RuntimeError):
+        await run.retry(func, on_failed=lambda args: events.append(f'failed {args.attempt}/{args.max_attempts}'))
+    assert events == ['call', 'failed 0/3',
+                      'call', 'failed 1/3',
+                      'call', 'failed 2/3']