File size: 3,700 Bytes
e816bb2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import asyncio
import functools
import inspect
from collections.abc import Callable
from ..exceptions import APIError
DELAY_FACTOR = 5
def running(retry: int = 0) -> Callable:
"""
Decorator to check if GeminiClient is running before making a request.
Supports both regular async functions and async generators.
Parameters
----------
retry: `int`, optional
Max number of retries when `gemini_webapi.APIError` is raised.
"""
def decorator(func):
if inspect.isasyncgenfunction(func):
@functools.wraps(func)
async def wrapper(client, *args, current_retry=None, **kwargs):
if current_retry is None:
current_retry = retry
try:
if not client._running:
await client.init(
timeout=client.timeout,
auto_close=client.auto_close,
close_delay=client.close_delay,
auto_refresh=client.auto_refresh,
refresh_interval=client.refresh_interval,
verbose=client.verbose,
watchdog_timeout=client.watchdog_timeout,
)
if not client._running:
raise APIError(
f"Invalid function call: GeminiClient.{func.__name__}. Client initialization failed."
)
async for item in func(client, *args, **kwargs):
yield item
except APIError:
if current_retry > 0:
delay = (retry - current_retry + 1) * DELAY_FACTOR
await asyncio.sleep(delay)
async for item in wrapper(
client, *args, current_retry=current_retry - 1, **kwargs
):
yield item
else:
raise
return wrapper
else:
@functools.wraps(func)
async def wrapper(client, *args, current_retry=None, **kwargs):
if current_retry is None:
current_retry = retry
try:
if not client._running:
await client.init(
timeout=client.timeout,
auto_close=client.auto_close,
close_delay=client.close_delay,
auto_refresh=client.auto_refresh,
refresh_interval=client.refresh_interval,
verbose=client.verbose,
watchdog_timeout=client.watchdog_timeout,
)
if not client._running:
raise APIError(
f"Invalid function call: GeminiClient.{func.__name__}. Client initialization failed."
)
return await func(client, *args, **kwargs)
except APIError:
if current_retry > 0:
delay = (retry - current_retry + 1) * DELAY_FACTOR
await asyncio.sleep(delay)
return await wrapper(
client, *args, current_retry=current_retry - 1, **kwargs
)
raise
return wrapper
return decorator
|