feat: add on_connected and on_disconnected callback

Author: Undertone0809

README.md

@@ -41,32 +41,59 @@ Here are some minimal example programs using the `cushy-socket`: a server that e
 
 ```python
 # echo tcp server program
+import socket
 from cushy_socket.tcp import CushyTCPServer
 
-es_tcp_server = CushyTCPServer(host='localhost', port=7777)
-es_tcp_server.run()
+cushy_tcp_server = CushyTCPServer(host='localhost', port=7777)
+cushy_tcp_server.run()
 
 
-@es_tcp_server.on_message()
+@cushy_tcp_server.on_connected()
+def handle_on_connected(sock: socket.socket):
+    print(f"[server decorator callback] new client connected.")
+    print(sock)
+
+
+@cushy_tcp_server.on_disconnected()
+def handle_on_disconnected(sock: socket.socket):
+    print(f"[server decorator callback] a client disconnected.")
+    print(sock)
+
+
+@cushy_tcp_server.on_message()
 def handle_msg_from_client(msg: str):
-    print(f"[server decorator callback] es_tcp_server rec msg: {msg}")
-    es_tcp_server.send("hello, I am server")
+    print(f"[server decorator callback] cushy_tcp_server rec msg: {msg}")
+    cushy_tcp_server.send("hello, I am server")
+
+
 ```
 
 ```python
 # echo tcp client program
 from cushy_socket.tcp import CushyTCPClient
 
-es_tcp_client = CushyTCPClient(host='localhost', port=7777)
-es_tcp_client.run()
+cushy_tcp_client = CushyTCPClient(host='localhost', port=7777)
+cushy_tcp_client.run()
 
 
-@es_tcp_client.on_message()
+@cushy_tcp_client.on_connected()
+def handle_on_connected():
+    print(f"[client decorator callback] connect to server.")
+
+
+@cushy_tcp_client.on_disconnected()
+def handle_on_disconnected():
+    print(f"[client decorator callback] server disconnected.")
+
+
+@cushy_tcp_client.on_message()
 def handle_msg_from_server(msg: str):
-    print(f"[client decorator callback] es_tcp_client rec msg: {msg}")
+    print(f"[client decorator callback] cushy_tcp_client rec msg: {msg}")
+
 
+cushy_tcp_client.send("hello, here is CSTCP client")
+cushy_tcp_client.close()
 
-es_tcp_client.send("hello, here is CSTCP client")
 ```
 
 

cushy_socket/tcp.py

@@ -20,31 +20,57 @@
 
 from threading import Thread
 from concurrent.futures import ThreadPoolExecutor
-from typing import List, Callable
+from typing import List, Callable, Optional
 import socket
 import logging
 
 __all__ = ['CushyTCPClient', 'CushyTCPServer']
-logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+enable_log = False
+if enable_log:
+    logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
 
 class CushyTCPClient:
     def __init__(self, host: str, port: int):
-        self.logger = logging.getLogger(__name__)
+        self.logger = logger
         self.host = host
         self.port = port
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.callbacks: List[Callable] = []
         self.is_running = False
+        self.executor = ThreadPoolExecutor()
+        self._callbacks: List[Callable] = []
+        self._disconnected_callback: Optional[Callable] = None
+        self._connected_callback: Optional[Callable] = None
 
     def run(self):
         """
         startup CSTCP Client
         """
         self.sock.connect((self.host, self.port))
         self.is_running = True
+        if self._connected_callback:
+            self.executor.submit(self._connected_callback)
         Thread(target=self._recv_thread).start()
 
+    def _recv_thread(self):
+        while True:
+            try:
+                msg = self.sock.recv(1024).decode('utf-8')
+            except Exception as e:
+                self.logger.error(f"[easy-socket] Error when receiving msg from server: {e}")
+                break
+            if not msg:
+                self.logger.error("[easy-socket] Server connection closed.")
+                break
+            self.logger.debug(f"[easy-socket] Received msg from server: {msg}")
+            for callback in self._callbacks:
+                self.executor.submit(callback, msg)
+
+        if self._disconnected_callback:
+            self.executor.submit(self._disconnected_callback)
+
     def send(self, msg: str or bytes):
         if type(msg) == str:
             self._send(msg.encode('utf-8'))
@@ -56,19 +82,8 @@ def send(self, msg: str or bytes):
     def _send(self, msg: bytes):
         self.sock.sendall(msg)
 
-    def _recv_thread(self):
-        while True:
-            msg = self.sock.recv(1024).decode('utf-8')
-            if not msg:
-                self.logger.error("[cushy-socket] Server connection closed.")
-                break
-            self.logger.info(f"[cushy-socket] Received msg from server: {msg}")
-
-            for callback in self.callbacks:
-                callback(msg)
-
     def listen(self, callback: Callable):
-        self.callbacks.append(callback)
+        self._callbacks.append(callback)
 
     def on_message(self):
         """
@@ -79,37 +94,54 @@ def on_message(self):
         ----------------------------------------------------------------------
         from cushy_socket.tcp import CushyTCPClient
 
-        es_tcp_client = CushyTCPClient(host='localhost', port=7777)
-        es_tcp_client.run()
+        cushy_tcp_client = CushyTCPClient(host='localhost', port=7777)
+        cushy_tcp_client.run()
 
 
-        @es_tcp_client.on_message()
+        @cushy_tcp_client.on_message()
         def handle_msg_from_server(msg: str):
-            print(f"[client decorator callback] es_tcp_client rec msg: {msg}")
+            print(f"[client decorator callback] cushy_tcp_client rec msg: {msg}")
         ----------------------------------------------------------------------
         """
+
+        def decorator(func):
+            self._callbacks.append(func)
+            return func
+
+        return decorator
+
+    def on_connected(self):
+        def decorator(func):
+            self._connected_callback = func
+            return func
+
+        return decorator
+
+    def on_disconnected(self):
         def decorator(func):
-            self.callbacks.append(func)
+            self._disconnected_callback = func
             return func
 
         return decorator
 
     def close(self):
+        self.sock.shutdown(2)
         self.sock.close()
         self.is_running = False
 
 
 class CushyTCPServer:
     def __init__(self, host: str, port: int):
-        self.logger = logging.getLogger(__name__)
+        self.logger = logger
         self.host = host
         self.port = port
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        # self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         self.clients = set()
-        self.callbacks: List[Callable] = []
         self.is_running = False
         self.executor = ThreadPoolExecutor()
+        self._callbacks: List[Callable] = []
+        self._disconnected_callback: Optional[Callable] = None
+        self._connected_callback: Optional[Callable] = None
 
     def run(self):
         """
@@ -118,14 +150,15 @@ def run(self):
         self.sock.bind((self.host, self.port))
         self.sock.listen()
         self.is_running = True
-        # self.executor.submit(self._accept_thread)
         Thread(target=self._accept_thread).start()
 
     def _accept_thread(self):
         while True:
             client_sock, client_addr = self.sock.accept()
-            self.logger.info(f"[cushy-socket] New client connected: {client_addr}")
+            self.logger.debug(f"[cushy-socket] New client connected: {client_addr}")
             self.clients.add(client_sock)
+            if self._connected_callback:
+                self.executor.submit(self._connected_callback, client_sock)
             self.executor.submit(self._recv_thread, client_sock)
 
     def _recv_thread(self, sock: socket.socket):
@@ -138,18 +171,21 @@ def _recv_thread(self, sock: socket.socket):
             try:
                 msg = sock.recv(1024).decode('utf-8')
             except Exception as e:
-                self.logger.error(f"[cushy-socket] Error when receiving msg from client: {e}")
-                self.clients.remove(sock)
-                sock.close()
+                self._client_close(sock, 'error', f"[cushy-socket] Error when receiving msg from client: {e}")
                 break
             if not msg:
-                self.logger.info("[cushy-socket] Client connection closed.")
-                self.clients.remove(sock)
-                sock.close()
+                self._client_close(sock, 'info', "[cushy-socket] Client connection closed.")
                 break
-            self.logger.info(f"[cushy-socket] Received msg from client: {msg}")
+            self.logger.debug(f"[cushy-socket] Received msg from client: {msg}")
             self.executor.submit(self._callback_thread, msg)
 
+    def _client_close(self, sock: socket.socket, log_type: str, log_msg: str):
+        self.logger.debug(log_msg) if log_type == 'info' else self.logger.debug(log_msg)
+        if self._disconnected_callback:
+            self.executor.submit(self._disconnected_callback, sock)
+        self.clients.remove(sock)
+        sock.close()
+
     def send(self, msg: str or bytes, sock: socket.socket = None):
         """
         send message to connected socket. You can choose specify socket client send message
@@ -169,11 +205,11 @@ def _send(self, msg: str or bytes, sock: socket.socket):
             raise Exception("Incorrect data type")
 
     def _callback_thread(self, msg: str):
-        for callback in self.callbacks:
-            callback(msg)
+        for callback in self._callbacks:
+            self.executor.submit(callback, msg)
 
     def listen(self, callback: Callable):
-        self.callbacks.append(callback)
+        self._callbacks.append(callback)
 
     def on_message(self):
         """
@@ -184,18 +220,33 @@ def on_message(self):
         ----------------------------------------------------------------------
         from cushy_socket.tcp import CushyTCPServer
 
-        es_tcp_server = CushyTCPServer(host='localhost', port=7777)
-        es_tcp_server.run()
+        cushy_tcp_server = CushyTCPServer(host='localhost', port=7777)
+        cushy_tcp_server.run()
 
 
-        @es_tcp_server.on_message()
+        @cushy_tcp_server.on_message()
         def handle_msg_from_client(msg: str):
-            print(f"[server decorator callback] es_tcp_server rec msg: {msg}")
-            es_tcp_server.send("hello, I am server")
+            print(f"[server decorator callback] cushy_tcp_server rec msg: {msg}")
+            cushy_tcp_server.send("hello, I am server")
         ----------------------------------------------------------------------
         """
+
+        def decorator(func):
+            self._callbacks.append(func)
+            return func
+
+        return decorator
+
+    def on_connected(self):
+        def decorator(func):
+            self._connected_callback = func
+            return func
+
+        return decorator
+
+    def on_disconnected(self):
         def decorator(func):
-            self.callbacks.append(func)
+            self._disconnected_callback = func
             return func
 
         return decorator

example/demo1_tcp_client.py

@@ -19,13 +19,24 @@
 
 from cushy_socket.tcp import CushyTCPClient
 
-es_tcp_client = CushyTCPClient(host='localhost', port=7777)
-es_tcp_client.run()
+cushy_tcp_client = CushyTCPClient(host='localhost', port=7777)
+cushy_tcp_client.run()
 
 
-@es_tcp_client.on_message()
+@cushy_tcp_client.on_connected()
+def handle_on_connected():
+    print(f"[client decorator callback] connect to server.")
+
+
+@cushy_tcp_client.on_disconnected()
+def handle_on_disconnected():
+    print(f"[client decorator callback] server disconnected.")
+
+
+@cushy_tcp_client.on_message()
 def handle_msg_from_server(msg: str):
-    print(f"[client decorator callback] es_tcp_client rec msg: {msg}")
+    print(f"[client decorator callback] cushy_tcp_client rec msg: {msg}")
 
 
-es_tcp_client.send("hello, here is CSTCP client")
+cushy_tcp_client.send("hello, here is CSTCP client")
+cushy_tcp_client.close()

example/demo1_tcp_server.py

@@ -17,13 +17,26 @@
 # Project Link: https://github.com/Undertone0809/cushy-socket
 # Contact Email: [email protected]
 
+import socket
 from cushy_socket.tcp import CushyTCPServer
 
-es_tcp_server = CushyTCPServer(host='localhost', port=7777)
-es_tcp_server.run()
+cushy_tcp_server = CushyTCPServer(host='localhost', port=7777)
+cushy_tcp_server.run()
 
 
-@es_tcp_server.on_message()
+@cushy_tcp_server.on_connected()
+def handle_on_connected(sock: socket.socket):
+    print(f"[server decorator callback] new client connected.")
+    print(sock)
+
+
+@cushy_tcp_server.on_disconnected()
+def handle_on_disconnected(sock: socket.socket):
+    print(f"[server decorator callback] a client disconnected.")
+    print(sock)
+
+
+@cushy_tcp_server.on_message()
 def handle_msg_from_client(msg: str):
-    print(f"[server decorator callback] es_tcp_server rec msg: {msg}")
-    es_tcp_server.send("hello, I am server")
+    print(f"[server decorator callback] cushy_tcp_server rec msg: {msg}")
+    cushy_tcp_server.send("hello, I am server")

setup.py

@@ -20,7 +20,7 @@
 
 setuptools.setup(
     name="cushy-socket",
-    version="1.1.0",
+    version="1.2.0",
     author="Zeeland",
     author_email="[email protected]",
     description="A lightweight socket library",