Add decorator for websocket server events

server.py

@@ -1,26 +1,29 @@
 from websocket_server import WebsocketServer
 
+
+PORT=9001
+server = WebsocketServer(port = PORT)
+
+
 # Called for every client connecting (after handshake)
+@server.on_new_client()
 def new_client(client, server):
 	print("New client connected and was given id %d" % client['id'])
 	server.send_message_to_all("Hey all, a new client has joined us")
 
 
 # Called for every client disconnecting
+@server.on_client_left()
 def client_left(client, server):
 	print("Client(%d) disconnected" % client['id'])
 
 
 # Called when a client sends a message
+@server.on_message_received()
 def message_received(client, server, message):
 	if len(message) > 200:
 		message = message[:200]+'..'
 	print("Client(%d) said: %s" % (client['id'], message))
 
 
-PORT=9001
-server = WebsocketServer(port = PORT)
-server.set_fn_new_client(new_client)
-server.set_fn_client_left(client_left)
-server.set_fn_message_received(message_received)
 server.run_forever()

websocket_server/websocket_server.py

@@ -14,6 +14,8 @@
 
 from websocket_server.thread import WebsocketServerThread
 
+from typing import Callable
+
 logger = logging.getLogger(__name__)
 logging.basicConfig()
 
@@ -52,7 +54,6 @@
 
 
 class API():
-
     def run_forever(self, threaded=False):
         return self._run_forever(threaded)
 
@@ -64,14 +65,32 @@ def client_left(self, client, server):
 
     def message_received(self, client, server, message):
         pass
-
-    def set_fn_new_client(self, fn):
+
+    def on_new_client(self):
+        def decorator(func: Callable):
+            self.set_fn_new_client(func)
+
+        return decorator
+
+    def on_client_left(self):
+        def decorator(func: Callable):
+            self.set_fn_client_left(func)
+
+        return decorator
+
+    def on_message_received(self):
+        def decorator(func: Callable):
+            self.set_fn_message_received(func)
+
+        return decorator
+
+    def set_fn_new_client(self, fn: Callable):
         self.new_client = fn
 
-    def set_fn_client_left(self, fn):
+    def set_fn_client_left(self, fn: Callable):
         self.client_left = fn
 
-    def set_fn_message_received(self, fn):
+    def set_fn_message_received(self, fn: Callable):
         self.message_received = fn
 
     def send_message(self, client, msg):