From 905dccc243aadab28c79081cbcd28883ab1309ca Mon Sep 17 00:00:00 2001
From: ktianc <ktianc2000@163.com>
Date: Fri, 5 Jan 2024 11:33:35 +0800
Subject: [PATCH] =?UTF-8?q?socket=20client=20=E6=B7=BB=E5=8A=A0=20tcp=20?=
 =?UTF-8?q?=E9=80=9A=E4=BF=A1=E5=8D=8F=E8=AE=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 kinit-api/utils/socket_client.py | 59 ++++++++++++++++++++++++++++----
 1 file changed, 52 insertions(+), 7 deletions(-)

diff --git a/kinit-api/utils/socket_client.py b/kinit-api/utils/socket_client.py
index efc1426..0f53a92 100644
--- a/kinit-api/utils/socket_client.py
+++ b/kinit-api/utils/socket_client.py
@@ -7,23 +7,65 @@ class SocketClient:
     socket 客户端操作
     """
 
-    def __init__(self, host: str = "127.0.0.1", port: int = 3636):
+    def __init__(self, host: str = "127.0.0.1", port: int = 3636, send_type: str = "tcp"):
         """
-        :param host:
-        :param port:
+        :param host: socket server 地址
+        :param port: socket server 端口
+        :param send_type: 通信协议
         """
-        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.send_type = send_type
+        if self.send_type == "tcp":
+            socket_type = socket.SOCK_STREAM
+        elif self.send_type == "udp":
+            socket_type = socket.SOCK_DGRAM
+        else:
+            print("不支持通信协议")
+            raise ValueError("不支持的通信协议")
+        self.client_socket = socket.socket(socket.AF_INET, socket_type)
         self.host = host
         self.port = port
+        if self.send_type == "tcp":
+            self.tcp_connect()
+
+    def tcp_connect(self):
+        """
+        TCP 连接服务端
+        :return:
+        """
+        self.client_socket.connect((self.host, self.port))
+        print("tcp 连接成功")
 
     def udp_send_message(self, message: str):
         """
-        发送消息
+        UDP 发送消息
         :param message:
         :return:
         """
-        # 如果你想接收响应
         self.client_socket.sendto(message.encode('utf-8'), (self.host, self.port))
+        print("udp 消息发送成功：", message)
+
+    def tcp_send_message(self, message: str):
+        """
+        TCP 发送消息
+        :param message:
+        :return:
+        """
+        self.client_socket.sendall(message.encode('utf-8'))
+        print("tcp 消息发送成功：", message)
+
+    def send_message(self, message: str):
+        """
+        TCP 发送消息
+        :param message:
+        :return:
+        """
+        if self.send_type == "tcp":
+            self.tcp_send_message(message)
+        elif self.send_type == "udp":
+            self.udp_send_message(message)
+        else:
+            print("不支持协议")
+            raise ValueError("不支持的协议")
 
     def close(self):
         """
@@ -34,6 +76,9 @@ class SocketClient:
 
 
 if __name__ == '__main__':
+    _host = "127.0.0.1"
+    _port = 3636
+
     SC = SocketClient()
-    SC.udp_send_message(json.dumps({"training_id": 1, "instruct": "4"}))
+    SC.tcp_send_message(json.dumps({"label": "ceshi", "value": 1}))
     SC.close()