Allow custom typing useWebSocket to get JsonObject type

robtaussig · May 20, 2022 · 7f70687 · 7f70687
1 parent b8fba3b
commit 7f70687
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/src/lib/types.ts b/src/lib/types.ts
@@ -51,11 +51,11 @@ export type Subscriber<T = WebSocketEventMap['message']> = {
   reconnect: MutableRefObject<() => void>,
 }
 
-export type WebSocketHook<T = WebSocketEventMap['message'], P = JsonValue> = {
+export type WebSocketHook<T = JsonValue, P = WebSocketEventMap['message'] | null> = {
   sendMessage: SendMessage,
   sendJsonMessage: SendJsonMessage,
-  lastMessage: T | null,
-  lastJsonMessage: P,
+  lastMessage: P,
+  lastJsonMessage: T,
   readyState: ReadyState,
   getWebSocket: () => (WebSocketLike | null),
 }

diff --git a/src/lib/use-websocket.ts b/src/lib/use-websocket.ts
@@ -12,25 +12,26 @@ import {
   WebSocketMessage,
   WebSocketHook,
   WebSocketLike,
+  JsonValue,
 } from './types';
 import { assertIsWebSocket } from './util';
 
-export const useWebSocket = (
+export const useWebSocket = <T extends JsonValue | null = JsonValue | null>(
   url: string | (() => string | Promise<string>) | null,
   options: Options = DEFAULT_OPTIONS,
   connect: boolean = true,
-): WebSocketHook => {
+): WebSocketHook<T> => {
   const [lastMessage, setLastMessage] = useState<WebSocketEventMap['message'] | null>(null);
   const [readyState, setReadyState] = useState<ReadyStateState>({});
-  const lastJsonMessage = useMemo(() => {
+  const lastJsonMessage: T = useMemo(() => {
     if (lastMessage) {
       try {
         return JSON.parse(lastMessage.data);
       } catch (e) {
         return UNPARSABLE_JSON_OBJECT;
       }
     }
-    return null;
+    return null; 
   },[lastMessage]);
   const convertedUrl = useRef<string | null>(null);
   const webSocketRef = useRef<WebSocketLike | null>(null);