Support custom JSON encoders

wavelink/client.py

@@ -25,6 +25,7 @@
 import logging
 from discord.ext import commands
 from functools import partial
+from json import dumps
 from typing import Optional, Union
 
 from .errors import *
@@ -69,6 +70,8 @@ def __init__(self, bot: Union[commands.Bot, commands.AutoShardedBot]):
 
         self.nodes = {}
 
+        self._dumps = dumps
+
         bot.add_listener(self.update_handler, 'on_socket_response')
 
     @property
@@ -157,7 +160,7 @@ async def get_tracks(self, query: str, *, retry_on_failure: bool = True) -> Opti
             There are no :class:`wavelink.node.Node`s currently connected.
         """
         node = self.get_best_node()
-        
+
         if node is None:
             raise ZeroConnectedNodes
 
@@ -390,7 +393,7 @@ async def initiate_node(self, host: str, port: int, *, rest_uri: str, password:
             Whether the websocket should be started with the secure wss protocol.
         heartbeat: Optional[float]
             Send ping message every heartbeat seconds and wait pong response, if pong response is not received then close connection.
-        
+
         Returns
         ---------
         :class:`wavelink.node.Node`
@@ -416,8 +419,9 @@ async def initiate_node(self, host: str, port: int, *, rest_uri: str, password:
                     session=self.session,
                     client=self,
                     secure=secure,
-                    heartbeat=heartbeat)
-
+                    heartbeat=heartbeat,
+                    dumps=self._dumps)
+
         await node.connect(bot=self.bot)
 
         node.available = True
@@ -471,3 +475,18 @@ async def update_handler(self, data) -> None:
                 pass
             else:
                 await player._voice_state_update(data['d'])
+
+    def set_serializer(self, serializer_function) -> None:
+        """Sets the JSON dumps function for use in the websocket.
+        The default one is the built-in JSON module.
+
+        Parameters
+        ----------
+        serializer_function: Callable[[Dict[str, Any]]], Union[str, bytes]]
+            The function that serializes the JSON data to a string or bytes.
+        """
+        self._dumps = serializer_function
+        # Update all existing nodes
+        for node in self.nodes.values():
+            node._dumps = serializer_function
+            node._websocket._dumps = serializer_function

wavelink/node.py

@@ -22,9 +22,10 @@
 """
 import asyncio
 import inspect
+import json
 import logging
 from discord.ext import commands
-from typing import Optional, Union
+from typing import Any, Callable, Dict, Optional, Union
 from urllib.parse import quote
 
 from .backoff import ExponentialBackoff
@@ -69,7 +70,8 @@ def __init__(self, host: str,
                  identifier: str,
                  shard_id: int = None,
                  secure: bool = False,
-                 heartbeat: float = None
+                 heartbeat: float = None,
+                 dumps: Callable[[Dict[str, Any]], Union[str, bytes]] = json.dumps
                  ):
 
         self.host = host
@@ -83,6 +85,8 @@ def __init__(self, host: str,
         self.secure = secure
         self.heartbeat = heartbeat
 
+        self._dumps = dumps
+
         self.shard_id = shard_id
 
         self.players = {}
@@ -127,7 +131,8 @@ async def connect(self, bot: Union[commands.Bot, commands.AutoShardedBot]) -> No
                                     password=self.password,
                                     shard_count=self.shards,
                                     user_id=self.uid,
-                                    secure=self.secure)
+                                    secure=self.secure,
+                                    dumps=self._dumps)
         await self._websocket._connect()
 
         __log__.info(f'NODE | {self.identifier} connected:: {self.__repr__()}')

wavelink/websocket.py

@@ -47,6 +47,7 @@ def __init__(self, **attrs):
         self.shard_count = attrs.get('shard_count')
         self.user_id = attrs.get('user_id')
         self.secure = attrs.get('secure')
+        self._dumps = attrs.get('dumps')
 
         self._websocket = None
         self._last_exc = None
@@ -165,4 +166,12 @@ def _get_event_payload(self, name: str, data):
     async def _send(self, **data):
         if self.is_connected:
             __log__.debug(f'WEBSOCKET | Sending Payload:: {data}')
-            await self._websocket.send_json(data)
+            data_str = self._dumps(data)
+            if isinstance(data_str, bytes):
+                # Some JSON libraries serialize to bytes
+                # Yet Lavalink does not support binary websockets
+                # So we need to decode. In the future, maybe
+                # self._websocket.send_bytes could be used
+                # if Lavalink ever implements it
+                data_str = data_str.decode('utf-8')
+            await self._websocket.send_str(data_str)