From 87b8a55625ac2bca306b25ff197d9955dfb787ec Mon Sep 17 00:00:00 2001
From: Lev Gorodetskiy <github@droserasprout.space>
Date: Fri, 2 Jun 2023 15:54:31 -0300
Subject: [PATCH] `BaseJSONProtocol` and `BaseWebsocketTransport` classes for
 plain JSON over WebSockets and custom protocols

---
 CHANGELOG.md                         |  6 ++++++
 src/pysignalr/messages.py            |  9 +++++++++
 src/pysignalr/protocol/json.py       | 17 +++++++++++++++++
 src/pysignalr/transport/websocket.py |  8 ++++++++
 4 files changed, 40 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6a8717d..9f24936 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog], and this project adheres to [Semantic Versioning].
 
+## [Unreleased]
+
+### Added
+
+- `BaseJSONProtocol` and `BaseWebsocketTransport` classes for plain JSON over WebSockets and custom protocols.
+
 ## [0.2.0] - 2023-04-07
 
 ### Added
diff --git a/src/pysignalr/messages.py b/src/pysignalr/messages.py
index 8acbb6e..0deaa5c 100644
--- a/src/pysignalr/messages.py
+++ b/src/pysignalr/messages.py
@@ -328,3 +328,12 @@ class StreamItemMessage(Message, type_=MessageType.stream_item):
     invocation_id: str
     item: Any
     headers: Optional[Dict[str, Any]] = None
+
+
+class JSONMessage(Message, type_=MessageType._):
+    """Not a real message type; used in BaseJSONProtocol to skip pysignalr-specific things"""
+    def __init__(self, data: Dict[str, Any]) -> None:
+        self.data = data
+
+    def dump(self) -> Dict[str, Any]:
+        return self.data
diff --git a/src/pysignalr/protocol/json.py b/src/pysignalr/protocol/json.py
index 59ea5a7..6aad980 100644
--- a/src/pysignalr/protocol/json.py
+++ b/src/pysignalr/protocol/json.py
@@ -11,10 +11,12 @@
 from pysignalr.messages import CloseMessage  # 7
 from pysignalr.messages import CompletionMessage  # 3
 from pysignalr.messages import HandshakeMessage
+from pysignalr.messages import HandshakeRequestMessage
 from pysignalr.messages import HandshakeResponseMessage
 from pysignalr.messages import InvocationMessage  # 1
 from pysignalr.messages import Message
 from pysignalr.messages import MessageType
+from pysignalr.messages import JSONMessage  # virtual
 from pysignalr.messages import PingMessage  # 6
 from pysignalr.messages import StreamInvocationMessage  # 4
 from pysignalr.messages import StreamItemMessage  # 2
@@ -31,6 +33,21 @@ def default(self, obj: Union[Message, MessageType]) -> Union[str, int, Dict[str,
 message_encoder = MessageEncoder()
 
 
+class BaseJSONProtocol(Protocol):
+    def __init__(self) -> None:
+        pass
+
+    def decode(self, raw_message: Union[str, bytes]) -> Tuple[JSONMessage]:
+        json_message = orjson.loads(raw_message)
+        return (JSONMessage(data=json_message),)
+
+    def encode(self, message: Union[Message, HandshakeRequestMessage]) -> Union[str, bytes]:
+        return orjson.dumps(message.dump())
+
+    def decode_handshake(self, raw_message: Union[str, bytes]) -> Tuple[HandshakeResponseMessage, Iterable[Message]]:
+        raise NotImplementedError
+
+
 class JSONProtocol(Protocol):
     def __init__(self) -> None:
         super().__init__(
diff --git a/src/pysignalr/transport/websocket.py b/src/pysignalr/transport/websocket.py
index 79d72c7..7cb8319 100644
--- a/src/pysignalr/transport/websocket.py
+++ b/src/pysignalr/transport/websocket.py
@@ -212,3 +212,11 @@ async def _on_raw_message(self, raw_message: Union[str, bytes]) -> None:
 
     async def _on_message(self, message: Message) -> None:
         await self._callback(message)
+
+
+class BaseWebsocketTransport(WebsocketTransport):
+    async def _keepalive(self, conn: WebSocketClientProtocol) -> None:
+        return
+
+    async def _handshake(self, conn: WebSocketClientProtocol) -> None:
+        return