Fix retry_on handling in run_with_retry

- Correctly distinguish between exception classes, lists/tuples of exception classes, and callables.
- Add support for lists in `retry_on`, alongside tuples.
- Prevent exception classes from being incorrectly treated as callables.
- Raise a `TypeError` if `retry_on` is of an unsupported type.
This commit is contained in:
Itay Etelis
2024-09-18 13:41:22 +03:00
parent 6873229bbf
commit adf5cb0ff1
+22 -6
View File
@@ -38,11 +38,19 @@ def run_with_retry(
# increment attempts
attempts += 1
# check if we should retry
if callable(retry_policy.retry_on):
if isinstance(retry_policy.retry_on, (list, tuple)):
if not isinstance(exc, tuple(retry_policy.retry_on)):
raise
elif isinstance(retry_policy.retry_on, type) and issubclass(retry_policy.retry_on, Exception):
if not isinstance(exc, retry_policy.retry_on):
raise
elif callable(retry_policy.retry_on):
if not retry_policy.retry_on(exc):
raise
elif not isinstance(exc, retry_policy.retry_on):
raise
else:
raise TypeError(
"retry_on must be an Exception class, a list or tuple of Exception classes, or a callable"
)
# check if we should give up
if attempts >= retry_policy.max_attempts:
raise
@@ -94,11 +102,19 @@ async def arun_with_retry(
# increment attempts
attempts += 1
# check if we should retry
if callable(retry_policy.retry_on):
if isinstance(retry_policy.retry_on, (list, tuple)):
if not isinstance(exc, tuple(retry_policy.retry_on)):
raise
elif isinstance(retry_policy.retry_on, type) and issubclass(retry_policy.retry_on, Exception):
if not isinstance(exc, retry_policy.retry_on):
raise
elif callable(retry_policy.retry_on):
if not retry_policy.retry_on(exc):
raise
elif not isinstance(exc, retry_policy.retry_on):
raise
else:
raise TypeError(
"retry_on must be an Exception class, a list or tuple of Exception classes, or a callable"
)
# check if we should give up
if attempts >= retry_policy.max_attempts:
raise