22
33
44class AsyncRateLimiter :
5- def __init__ (self , storage_file = 'rate_limit_data .pkl' ):
5+ def __init__ (self , storage_file = "swibots .pkl" ):
66 self .storage_file = storage_file
77 self .rate_limit_data = {}
8-
8+
99 async def _load_data (self ):
1010 if os .path .exists (self .storage_file ):
11- async with aiofiles .open (self .storage_file , 'rb' ) as f :
11+ async with aiofiles .open (self .storage_file , "rb" ) as f :
1212 content = await f .read ()
1313 self .rate_limit_data = pickle .loads (content )
14-
14+
1515 async def _save_data (self ):
16- async with aiofiles .open (self .storage_file , 'wb' ) as f :
16+ async with aiofiles .open (self .storage_file , "wb" ) as f :
1717 await f .write (pickle .dumps (self .rate_limit_data ))
18-
18+
1919 def limit (self , key : str , times : int , seconds : int ):
2020 def decorator (func ):
2121 async def async_wrapper (* args , ** kwargs ):
2222 current_time = time .time ()
2323
2424 await self ._load_data ()
25-
25+
2626 if key not in self .rate_limit_data :
2727 self .rate_limit_data [key ] = []
28-
29- self .rate_limit_data [key ] = [t for t in self .rate_limit_data [key ] if current_time - t < seconds ]
30-
28+
29+ self .rate_limit_data [key ] = [
30+ t for t in self .rate_limit_data [key ] if current_time - t < seconds
31+ ]
32+
3133 if len (self .rate_limit_data [key ]) >= times :
32- raise Exception (f"Rate limit exceeded: { times } calls in { seconds } seconds" )
33-
34+ raise RateLimitException (
35+ f"Rate limit exceeded: { times } calls in { seconds } seconds"
36+ )
37+
3438 # Record the call
3539 self .rate_limit_data [key ].append (current_time )
3640 await self ._save_data ()
3741 return await func (* args , ** kwargs )
38-
42+
3943 return async_wrapper
44+
4045 return decorator
4146
4247
4348class RateLimitException (Exception ):
44- pass
49+ pass
0 commit comments