| Total Complexity | 4 |
| Total Lines | 41 |
| Duplicated Lines | 0 % |
| Changes | 0 | ||
| 1 | # Copyright Pincer 2021-Present |
||
|
|
|||
| 2 | # Full MIT License can be found in `LICENSE` at the project root. |
||
| 3 | |||
| 4 | from time import time |
||
| 5 | |||
| 6 | |||
| 7 | class SlidingWindow: |
||
| 8 | def __init__(self, capacity: int, time_unit: float): |
||
| 9 | self.capacity: int = capacity |
||
| 10 | self.time_unit: float = time_unit |
||
| 11 | |||
| 12 | self.__cur_time: float = time() |
||
| 13 | self.__pre_count: int = capacity |
||
| 14 | self.__cur_count: int = 0 |
||
| 15 | |||
| 16 | def allow(self) -> bool: |
||
| 17 | # TODO: fix docs |
||
| 18 | """ |
||
| 19 | |||
| 20 | Returns |
||
| 21 | ------- |
||
| 22 | |||
| 23 | """ |
||
| 24 | # Reset rate limit: |
||
| 25 | if (time() - self.__cur_time) > self.time_unit: |
||
| 26 | self.__cur_time = time() |
||
| 27 | self.__pre_count = self.__cur_count |
||
| 28 | self.__cur_count = 0 |
||
| 29 | |||
| 30 | # Calculate the estimated count |
||
| 31 | passed_time = time() - self.__cur_time |
||
| 32 | time_cnt = (self.time_unit - passed_time) / self.time_unit |
||
| 33 | est_cnt = self.__pre_count * time_cnt + self.__cur_count |
||
| 34 | |||
| 35 | # Request has passed the capacity |
||
| 36 | if est_cnt > self.capacity: |
||
| 37 | return False |
||
| 38 | |||
| 39 | self.__cur_count += 1 |
||
| 40 | return True |
||
| 41 |