feat: support backwards invocation

This commit is contained in:
Yeuoly
2024-07-29 18:57:34 +08:00
parent f29b44acd8
commit d52476c1c9
4 changed files with 77 additions and 52 deletions

View File

@ -1,7 +1,7 @@
import logging
import time
import uuid
from collections.abc import Generator
from collections.abc import Callable, Generator
from datetime import timedelta
from typing import Optional, Union
@ -91,7 +91,7 @@ class RateLimit:
class RateLimitGenerator:
def __init__(self, rate_limit: RateLimit, generator: Union[Generator, callable], request_id: str):
def __init__(self, rate_limit: RateLimit, generator: Union[Generator, Callable[[], Generator]], request_id: str):
self.rate_limit = rate_limit
if callable(generator):
self.generator = generator()