client.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. import socket
  2. from threading import Thread
  3. class SocketClient(Thread):
  4. __BLOCK_SIZE, __MAX_TIMES = 1024, 1024 * 10
  5. def __init__(self, host: "str", port: "int"):
  6. super().__init__()
  7. self.__client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  8. self.__client.connect((host, port))
  9. self.__handlers, self.__running = {"BROADCAST": self.__on_broad}, False
  10. def __on_broad(self, msg: "str"):
  11. raise NotImplementedError(
  12. "got a broadcast msg, but there are no handler for it.\n"
  13. "you should bind a <BROADCAST> message handler with 'on'."
  14. )
  15. def on(self, event: "str"):
  16. def decorator(func):
  17. self.__handlers[event] = func
  18. return func
  19. return decorator
  20. def emit(self, event: "str", msg: "bytes"):
  21. self.__client.sendall(f"==EVENT:{event}==".encode("utf-8") + msg + b"==END==")
  22. def run(self) -> None:
  23. self.__running = True
  24. try:
  25. while self.__running:
  26. times = SocketClient.__MAX_TIMES
  27. cur = self.__client.recv(SocketClient.__BLOCK_SIZE)
  28. if cur.startswith(b"==EVENT:"):
  29. _e = cur[8:].find(b"==")
  30. event, cur = cur[8:_e + 8].decode("utf-8"), cur[_e + 10:]
  31. data = cur
  32. while not data.endswith(b"==END=="):
  33. cur = self.__client.recv(SocketClient.__BLOCK_SIZE)
  34. data, times = data + cur, times - 1
  35. if times <= 0:
  36. break
  37. if data.endswith(b"==END=="):
  38. if event in self.__handlers.keys():
  39. self.__handlers[event](data[:-7])
  40. else:
  41. raise RuntimeError(f"bad ending in msg: {data[:20]}")
  42. finally:
  43. self.__client.close()
  44. def stop(self):
  45. self.__running = False