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 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,58 +63,33 @@ 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):
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: