async def __aenter__(self) -> "Gatherer": self.session_id = uuid.uuid4() self.session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False, limit=self.workers)) for pipeline in pipelines.values(): pipeline.sessions[self.session_id] = self.session LOGGER.log(self.log_level, f"[Trace: Pyot Gatherer] Created session '{self.session_id}'") return self
async def __aexit__(self, *args): await self.queue.join() for worker in self.workers: worker.cancel() await asyncio.gather(*self.workers, return_exceptions=True) self.is_joined = True LOGGER.log(self.log_level, f"[Trace: Pyot Queue] Joined {self.workers_num} workers") await self.session.close() LOGGER.log(self.log_level, f"[Trace: Pyot Queue] Closed session '{self.sid}'") for pipeline in pipelines.values(): pipeline.sessions.pop(self.sid) return
async def __aenter__(self) -> "Queue": self.queue = asyncio.Queue(maxsize=self.maxsize) self.responses = [] self.workers = [] for _ in range(self.workers_num): worker = asyncio.create_task(self.worker(self.queue)) self.workers.append(worker) LOGGER.log(self.log_level, f"[Trace: Pyot Queue] Spawned {self.workers_num} workers") self.sid = uuid.uuid4() LOGGER.log(self.log_level, f"[Trace: Pyot Queue] Created session '{self.sid}'") self.session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) for pipeline in pipelines.values(): pipeline.sessions[self.sid] = self.session return self
async def __aexit__(self, *args): await self.session.close() for pipeline in pipelines.values(): pipeline.sessions.pop(self.session_id) LOGGER.log(self.log_level, f"[Trace: Pyot Gatherer] Closed session '{self.session_id}'") return