feat: encode json and msgpack payloads

examples/simple_flask_example.py

@@ -18,7 +18,8 @@ def echo(ws: WebSocket) -> None:
 
     def publish_to_general() -> None:
         while 1:
-            websockets.publish(str(time.time()), ["time"])
+            websockets.publish({"time": time.time()}, ["time"], encode_json=True)
+            websockets.publish({"time": time.time()}, ["time"], encode_msgpack=True)
             time.sleep(1)
 
     Thread(target=publish_to_general).start()

pyproject.toml

@@ -8,6 +8,7 @@ authors = [
 dependencies = [
     "flask>=3.0.3",
     "wsproto>=1.2.0",
+    "msgspec>=0.18.6",
 ]
 requires-python = ">=3.10"
 readme = "README.md"

src/flask_websockets/websocket.py

@@ -10,6 +10,7 @@
 from time import time
 from typing import IO, TYPE_CHECKING, Any, cast
 
+import msgspec
 from wsproto import ConnectionType, WSConnection
 from wsproto.events import (
     AcceptConnection,
@@ -146,6 +147,22 @@ def send(self, data: bytes | str) -> None:
         except BrokenPipeError:
             self.state = WebSocketState.DISCONNECTED
 
+    def send_json(self, json_payload: Any) -> None:
+        """Sends json-encodable data to the websocket.
+
+        :param data: data to be sent to the websocket.
+        """
+        encoded_data = msgspec.json.encode(json_payload)
+        self.send(encoded_data)
+
+    def send_msgpack(self, msgspec_payload: Any) -> None:
+        """Sends msgpack-encodable data to the websocket.
+
+        :param data: data to be sent to the websocket.
+        """
+        encoded_data = msgspec.msgpack.encode(msgspec_payload)
+        self.send(encoded_data)
+
     def receive(self, timeout: int | float | None = None) -> bytes | str | None:
         """Receive data over the WebSocket connection.
 

src/flask_websockets/websockets.py

@@ -142,10 +142,14 @@ def unsubscribe(self, ws: WebSocket, channels: Iterable[str]) -> None:
             if channel in self._subscriptions:
                 self._subscriptions[channel].remove(ws)
 
-    def publish(self, data: bytes | str, channels: Iterable[str]) -> None:
+    def publish(
+        self, data: Any, channels: Iterable[str], encode_json: bool = False, encode_msgpack: bool = False
+    ) -> None:
         """Publish data to the given channels. All subscribed websockets will receive this data.
 
         :param data: Data to publish to the channels
+        :param encode_json: Whether data should be json-encoded before publishing
+        :param encode_msgpack: Whether data should be msgpack-encoded before publishing
         :param channels: Channels to publish the data to
         """
         channels_set: set[str] = set(channels)
@@ -155,6 +159,11 @@ def publish(self, data: bytes | str, channels: Iterable[str]) -> None:
 
             for ws in self._subscriptions[channel]:
                 try:
-                    ws.send(data)
+                    if encode_json:
+                        ws.send_json(data)
+                    elif encode_msgpack:
+                        ws.send_msgpack(data)
+                    else:
+                        ws.send(data)
                 except RuntimeError:
                     logger.warning("Couldn't to send data to a socket.")