diff --git a/src/snekbot/rpc.py b/src/snekbot/rpc.py index cb51028..67ae589 100644 --- a/src/snekbot/rpc.py +++ b/src/snekbot/rpc.py @@ -22,24 +22,20 @@ logger = logging.getLogger("snekbot.rpc") class RPC: class Response: def __init__(self, msg): - if isinstance(msg, list): - self.list = msg self.__dict__.update(msg) def __iter__(self): - for k in self.__dict__.get("data", []): - yield k + yield from self.__dict__.get("data", []) async def __aiter__(self): - for k in self.__dict__.get("data", []): - yield k + yield from self.__dict__.get("data", []) def __getitem__(self, name): - try: - return self.__dict__[name] - except: - pass - return self.__dict__.get("data", {})[name] + try: + return self.__dict__[name] + except: + pass + return self.__dict__.get("data", {})[name] def __setitem__(self, name, value): if name not in self.__dict__.get("data", {}): @@ -55,10 +51,6 @@ class RPC: self.current_call_id = None self.queue = asyncio.Queue() - async def queue(self, data): - logger.debug("Schedule for retry: " + str(data)) - await self.queue.put(data) - def __getattr__(self, name): async def method(*args, **kwargs): no_response = kwargs.pop("_no_response", False) @@ -71,59 +63,34 @@ class RPC: } await self.ws.send_json(payload) - async def returner(): + async def poller(): + while True: response = await self.ws.receive() data = response.json() - if data.get("callId") != self.current_call_id: - await self.queue.put(data) - continue - self.current_call_id = None - return self.Response(data) + if data.get("callId") == self.current_call_id: + self.current_call_id = None + return self.Response(data) + await self.queue.put(data) if no_response: return True - return await returner() + return await poller() return method - async def system(self, command): - if isinstance(command, str): - command = command.split(" ") - - path = pathlib.Path("output.txt") - - with path.open("w+") as f: - try: - subprocess.run(command, stderr=f, stdout=f) - except Exception as ex: - print("Error running command:", ex) - return f"Error: {ex}" - - response = None - - with path.open("r") as f: - response = f.read() - - try: - path.unlink() - except Exception as ex: - logger.error(ex) - - return response - - async def receive(self): - + async def receive(self): popped = [] while not self.queue.empty(): msg = await self.queue.get() if self.current_call_id == msg.get("callId"): - for m in popped: - await self.queue.put(m) self.current_call_id = None return self.Response(msg) - + popped.append(msg) + for m in popped: + await self.queue.put(m) + while True: try: msg = await self.ws.receive()