Update.

src/snek/research/serpentarium.py

@@ -33,76 +33,6 @@ class DatasetTable:
     def __getattr__(self, name):
         return DatasetMethod(self, name)
 
-class WebSocketClient:
-    def __init__(self):
-        self.buffer = b''
-        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.connect()     
-
-    def connect(self):
-        self.socket.connect(("127.0.0.1", 3131))
-        key = base64.b64encode(b'1234123412341234').decode('utf-8')
-        handshake = (
-            f"GET /db HTTP/1.1\r\n"
-            f"Host: localhost:3131\r\n"
-            f"Upgrade: websocket\r\n"
-            f"Connection: Upgrade\r\n"
-            f"Sec-WebSocket-Key: {key}\r\n"
-            f"Sec-WebSocket-Version: 13\r\n\r\n"
-        )
-        self.socket.sendall(handshake.encode('utf-8'))
-        response = self.read_until(b'\r\n\r\n')
-        if b'101 Switching Protocols' not in response:
-            raise Exception("Failed to connect to WebSocket")
-
-    def write(self, message):
-        message_bytes = message.encode('utf-8')
-        length = len(message_bytes)
-        if length <= 125:
-            self.socket.sendall(b'\x81' + bytes([length]) + message_bytes)
-        elif length >= 126 and length <= 65535:
-            self.socket.sendall(b'\x81' + bytes([126]) + length.to_bytes(2, 'big') + message_bytes)
-        else:
-            self.socket.sendall(b'\x81' + bytes([127]) + length.to_bytes(8, 'big') + message_bytes)
-        
-
-    def read_until(self, delimiter): 
-        while True:
-            find_pos = self.buffer.find(delimiter)
-            if find_pos != -1:
-                data = self.buffer[:find_pos+4]
-                self.buffer = self.buffer[find_pos+4:]
-                return data 
-            
-            chunk = self.socket.recv(1024)
-            if not chunk:
-                return None
-            self.buffer += chunk
-            
-    def read_exactly(self, length):
-        while len(self.buffer) < length:
-            chunk = self.socket.recv(length - len(self.buffer))
-            if not chunk:
-                return None
-            self.buffer += chunk 
-        response = self.buffer[: length]
-        self.buffer = self.buffer[length:]
-        return response
-
-    def read(self):
-        frame = None 
-        frame = self.read_exactly(2)
-        length = frame[1] & 127
-        if length == 126:
-            length = int.from_bytes(self.read_exactly(2), 'big')
-        elif length == 127:
-            length = int.from_bytes(self.read_exactly(8), 'big')
-        message = self.read_exactly(length)
-        return message
-    
-    def close(self):
-        self.socket.close()
-
 
 
 

src/snek/system/websocket.py

@@ -0,0 +1,81 @@
+class WebSocketClient:
+    def __init__(self, hostname, port):
+        self.buffer = b''
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.hostname = hostname 
+        self.port = port
+        self.connect()     
+    
+    def __getattr__(self, method, *args, **kwargs):
+        if method in self.__dict__.keys():
+            return self.__dict__[method]
+        def call(*args, **kwargs):
+            self.write(json.dumps({'method': method, 'args': args, 'kwargs': kwargs}))
+            return json.loads(self.read())
+        return call
+
+    def connect(self):
+        self.socket.connect((self.hostname, self.port))
+        key = base64.b64encode(b'1234123412341234').decode('utf-8')
+        handshake = (
+            f"GET /db HTTP/1.1\r\n"
+            f"Host: localhost:3131\r\n"
+            f"Upgrade: websocket\r\n"
+            f"Connection: Upgrade\r\n"
+            f"Sec-WebSocket-Key: {key}\r\n"
+            f"Sec-WebSocket-Version: 13\r\n\r\n"
+        )
+        self.socket.sendall(handshake.encode('utf-8'))
+        response = self.read_until(b'\r\n\r\n')
+        if b'101 Switching Protocols' not in response:
+            raise Exception("Failed to connect to WebSocket")
+
+    def write(self, message):
+        message_bytes = message.encode('utf-8')
+        length = len(message_bytes)
+        if length <= 125:
+            self.socket.sendall(b'\x81' + bytes([length]) + message_bytes)
+        elif length >= 126 and length <= 65535:
+            self.socket.sendall(b'\x81' + bytes([126]) + length.to_bytes(2, 'big') + message_bytes)
+        else:
+            self.socket.sendall(b'\x81' + bytes([127]) + length.to_bytes(8, 'big') + message_bytes)
+        
+
+    def read_until(self, delimiter): 
+        while True:
+            find_pos = self.buffer.find(delimiter)
+            if find_pos != -1:
+                data = self.buffer[:find_pos+4]
+                self.buffer = self.buffer[find_pos+4:]
+                return data 
+            
+            chunk = self.socket.recv(1024)
+            if not chunk:
+                return None
+            self.buffer += chunk
+            
+    def read_exactly(self, length):
+        while len(self.buffer) < length:
+            chunk = self.socket.recv(length - len(self.buffer))
+            if not chunk:
+                return None
+            self.buffer += chunk 
+        response = self.buffer[: length]
+        self.buffer = self.buffer[length:]
+        return response
+
+    def read(self):
+        frame = None 
+        frame = self.read_exactly(2)
+        length = frame[1] & 127
+        if length == 126:
+            length = int.from_bytes(self.read_exactly(2), 'big')
+        elif length == 127:
+            length = int.from_bytes(self.read_exactly(8), 'big')
+        message = self.read_exactly(length)
+        return message
+    
+    def close(self):
+        self.socket.close()
+
+