Update RPC.

This commit is contained in:
retoor 2025-08-02 02:59:33 +02:00
parent ce759a288f
commit dec19669cd

View File

@ -22,24 +22,20 @@ logger = logging.getLogger("snekbot.rpc")
class RPC: class RPC:
class Response: class Response:
def __init__(self, msg): def __init__(self, msg):
if isinstance(msg, list):
self.list = msg
self.__dict__.update(msg) self.__dict__.update(msg)
def __iter__(self): def __iter__(self):
for k in self.__dict__.get("data", []): yield from self.__dict__.get("data", [])
yield k
async def __aiter__(self): async def __aiter__(self):
for k in self.__dict__.get("data", []): yield from self.__dict__.get("data", [])
yield k
def __getitem__(self, name): def __getitem__(self, name):
try: try:
return self.__dict__[name] return self.__dict__[name]
except: except:
pass pass
return self.__dict__.get("data", {})[name] return self.__dict__.get("data", {})[name]
def __setitem__(self, name, value): def __setitem__(self, name, value):
if name not in self.__dict__.get("data", {}): if name not in self.__dict__.get("data", {}):
@ -55,10 +51,6 @@ class RPC:
self.current_call_id = None self.current_call_id = None
self.queue = asyncio.Queue() 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): def __getattr__(self, name):
async def method(*args, **kwargs): async def method(*args, **kwargs):
no_response = kwargs.pop("_no_response", False) no_response = kwargs.pop("_no_response", False)
@ -71,59 +63,34 @@ class RPC:
} }
await self.ws.send_json(payload) await self.ws.send_json(payload)
async def returner(): async def poller():
while True: while True:
response = await self.ws.receive() response = await self.ws.receive()
data = response.json() data = response.json()
if data.get("callId") != self.current_call_id: if data.get("callId") == self.current_call_id:
await self.queue.put(data) self.current_call_id = None
continue return self.Response(data)
self.current_call_id = None await self.queue.put(data)
return self.Response(data)
if no_response: if no_response:
return True return True
return await returner() return await poller()
return method return method
async def system(self, command): async def receive(self):
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):
popped = [] popped = []
while not self.queue.empty(): while not self.queue.empty():
msg = await self.queue.get() msg = await self.queue.get()
if self.current_call_id == msg.get("callId"): if self.current_call_id == msg.get("callId"):
for m in popped:
await self.queue.put(m)
self.current_call_id = None self.current_call_id = None
return self.Response(msg) return self.Response(msg)
popped.append(msg)
for m in popped:
await self.queue.put(m)
while True: while True:
try: try:
msg = await self.ws.receive() msg = await self.ws.receive()