Coverage for src/extratools_limit/rate_limit.py: 0%
45 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-22 18:25 -0700
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-22 18:25 -0700
1import asyncio
2import functools
3import random
4import time
5from collections.abc import Callable
6from datetime import timedelta
7from pathlib import Path
8from typing import Any
10from extratools_core.typing import PathLike
13class Wait:
14 def __init__(
15 self,
16 lockfile: PathLike | str,
17 *,
18 min_gap: timedelta | float = timedelta(seconds=0),
19 randomness: timedelta | float = timedelta(milliseconds=1),
20 use_async: bool = False,
21 ) -> None:
22 if isinstance(lockfile, str):
23 lockfile = Path(lockfile)
24 if isinstance(min_gap, timedelta):
25 min_gap = min_gap.seconds
26 if isinstance(randomness, timedelta):
27 randomness = randomness.seconds
29 self.__lockfile: PathLike = lockfile
30 self.__min_gap: float = min_gap
31 self.__randomness: float = randomness
32 self.__use_async: bool = use_async
34 def __call__(self, func: Callable[...]): # noqa: ANN204
35 @functools.wraps(func)
36 def wrapper(*args: Any, **kwargs: Any) -> Any:
37 if not self.__lockfile.is_file():
38 self.__lockfile.touch()
40 while True:
41 gap: float = time.time() - self.__lockfile.stat().st_mtime
42 if (remaining_gap := self.__min_gap - gap) > 0:
43 time.sleep(remaining_gap + random.random() * self.__randomness)
44 continue
46 # Note that since we are not actually locking the file,
47 # there is rare chance that multiple threads can run at the same time.
48 self.__lockfile.touch()
49 return func(*args, **kwargs)
51 @functools.wraps(func)
52 async def wrapper_async(*args: Any, **kwargs: Any) -> Any:
53 if not self.__lockfile.is_file():
54 self.__lockfile.touch()
56 while True:
57 gap: float = time.time() - self.__lockfile.stat().st_mtime
58 if (remaining_gap := self.__min_gap - gap) > 0:
59 await asyncio.sleep(remaining_gap + random.random() * self.__randomness)
60 continue
62 # Note that since we are not actually locking the file,
63 # there is rare chance that multiple threads can run at the same time.
64 self.__lockfile.touch()
65 return await func(*args, **kwargs)
67 return wrapper_async if self.__use_async else wrapper