Coverage for src/extratools_limit/rate_limit.py: 0%

1import asyncio 

2import functools 

3import random 

4import time 

5from datetime import timedelta 

6from pathlib import Path 

7from typing import Any 

8 

9from extratools_core.typing import PathLike 

10 

11 

12class Wait: 

13 def __init__( 

14 self, 

15 lockfile: PathLike | str, 

16 *, 

17 min_gap: timedelta | float = timedelta(seconds=0), 

18 randomness: timedelta | float = timedelta(milliseconds=1), 

19 use_async: bool = False, 

20 ) -> None: 

21 if isinstance(lockfile, str): 

22 lockfile = Path(lockfile) 

23 if isinstance(min_gap, timedelta): 

24 min_gap = min_gap.seconds 

25 if isinstance(randomness, timedelta): 

26 randomness = randomness.seconds 

27 

28 self.__lockfile: PathLike = lockfile 

29 self.__min_gap: float = min_gap 

30 self.__randomness: float = randomness 

31 self.__use_async: bool = use_async 

32 

33 def __call__(self, func): # noqa: ANN001, ANN204 

34 @functools.wraps(func) 

35 def wrapper(*args: Any, **kwargs: Any) -> Any: 

36 if not self.__lockfile.is_file(): 

37 self.__lockfile.touch() 

38 

39 while True: 

40 gap: float = time.time() - self.__lockfile.stat().st_mtime 

41 if (remaining_gap := self.__min_gap - gap) > 0: 

42 time.sleep(remaining_gap + random.random() * self.__randomness) 

43 continue 

44 

45 # Note that since we are not actually locking the file, 

46 # there is rare chance that multiple threads can run at the same time. 

47 self.__lockfile.touch() 

48 return func(*args, **kwargs) 

49 

50 @functools.wraps(func) 

51 async def wrapper_async(*args: Any, **kwargs: Any) -> Any: 

52 if not self.__lockfile.is_file(): 

53 self.__lockfile.touch() 

54 

55 while True: 

56 gap: float = time.time() - self.__lockfile.stat().st_mtime 

57 if (remaining_gap := self.__min_gap - gap) > 0: 

58 await asyncio.sleep(remaining_gap + random.random() * self.__randomness) 

59 continue 

60 

61 # Note that since we are not actually locking the file, 

62 # there is rare chance that multiple threads can run at the same time. 

63 self.__lockfile.touch() 

64 return await func(*args, **kwargs) 

65 

66 return wrapper_async if self.__use_async else wrapper