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

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 

9 

10from extratools_core.typing import PathLike 

11 

12 

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 

28 

29 self.__lockfile: PathLike = lockfile 

30 self.__min_gap: float = min_gap 

31 self.__randomness: float = randomness 

32 self.__use_async: bool = use_async 

33 

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() 

39 

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 

45 

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) 

50 

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() 

55 

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 

61 

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) 

66 

67 return wrapper_async if self.__use_async else wrapper