@@ -16,23 +16,80 @@ def init_logger(level=logging.INFO):
1616 logging .basicConfig (level = level , format = log_format )
1717
1818
19+ async def gather_with_concurrency (num_concurrent_tasks : int = None ,
20+ * tasks ):
21+ """
22+ Limit the number of tasks that are gather concurrently
23+ Args:
24+ num_concurrent_tasks(int): Number of concurrent tasks
25+ *tasks: Tasks to be gathered
26+
27+ Returns: result of tasks
28+
29+ """
30+ semaphore = asyncio .Semaphore (num_concurrent_tasks )
31+
32+ async def sem_task (task ):
33+ """
34+ If the num_concurrent_tasks is None, it does not \
35+ \
36+
37+ Args:
38+ task: task which might have to be throttled
39+
40+ Returns: result of task
41+
42+ """
43+ if num_concurrent_tasks is None :
44+ return await task
45+
46+ async with semaphore :
47+ return await task
48+
49+ all_tasks = (sem_task (task ) for task in tasks )
50+ return await asyncio .gather (* all_tasks )
51+
52+
1953async def gather_dict (tasks : dict ):
54+ """
55+ Gathers the tasks in a dict format (instead of typical List)
56+ Args:
57+ tasks(dict): Dictionary of tasks to be gathered.
58+ key represents the identifier, the value is the task
59+
60+ Returns: Result of the tasks in a dict format where key \
61+
62+
63+ """
2064 async def mark (key , coroutine ):
2165 return key , await coroutine
2266
67+ tasks_to_send = (mark (key , coroutine )
68+ for key , coroutine in tasks .items ())
69+
2370 return {
2471 key : result
25- for key , result in await asyncio . gather (
26- * ( mark ( key , coroutine ) for key , coroutine in tasks . items ())
27- )
72+ for key , result in await gather_with_concurrency ( 500 ,
73+ * tasks_to_send
74+ )
2875 }
2976
3077
31- class TokenBucket :
78+ class Borg :
79+ """Borg instance"""
80+ _shared_state = {}
81+
82+ def __init__ (self ):
83+ self .__dict__ = self ._shared_state
84+
85+
86+ class TokenBucket (Borg ):
3287 """Controls the number of requests that can be made to the API.
3388 All times are written in micro-seconds for a good level of accuracy"""
3489
35- def __init__ (self , request_limit : Dict , pause_seconds : float = 1 ):
90+ def __init__ (self ,
91+ request_limit : Dict ,
92+ pause_seconds : float = 1 ):
3693 """
3794 Initializes the TokenBucket algorithm to throttle the flow of requests
3895
@@ -42,16 +99,18 @@ def __init__(self, request_limit: Dict, pause_seconds: float = 1):
4299
43100 pause_seconds: amount of seconds to pause and test again
44101 """
45- self .queue = asyncio .Queue ()
46- (
47- self .bucket_list ,
48- self .delta_t_list ,
49- self .max_requests_list ,
50- ) = self ._initialize_buckets (request_limit )
51- self .last_check = datetime .now ()
52- self .pause_seconds = pause_seconds
53- self .results = []
54- self ._counter = 0
102+ super ().__init__ ()
103+ if not self ._shared_state :
104+ self .queue = asyncio .Queue ()
105+ (
106+ self .bucket_list ,
107+ self .delta_t_list ,
108+ self .max_requests_list ,
109+ ) = self ._initialize_buckets (request_limit )
110+ self .last_check = datetime .now ()
111+ self .pause_seconds = pause_seconds
112+ self .results = []
113+ self ._counter = 0
55114
56115 @staticmethod
57116 def _initialize_buckets (requests_dict : Dict ):
0 commit comments