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 |