Allow application to provide a custom JSON encoder/decoder.

miguelgrinberg · miguelgrinberg · commit f3967320883f · 2015-08-30T11:57:43.000-07:00
diff --git a/socketio/packet.py b/socketio/packet.py
@@ -1,5 +1,5 @@
 import functools
-import json
+import json as _json
 
 import six
 
@@ -11,6 +11,9 @@
 
 class Packet(object):
     """Socket.IO packet."""
+
+    json = _json
+
     def __init__(self, packet_type=EVENT, data=None, namespace=None, id=None,
                  binary=None, encoded_packet=None):
         self.packet_type = packet_type
@@ -54,7 +57,7 @@ def encode(self):
         if data is not None:
             if needs_comma:
                 encoded_packet += ','
-            encoded_packet += json.dumps(data, separators=(',', ':'))
+            encoded_packet += self.json.dumps(data, separators=(',', ':'))
         if attachments is not None:
             encoded_packet = [encoded_packet] + attachments
         return encoded_packet
@@ -90,7 +93,7 @@ def decode(self, encoded_packet):
                 self.id = self.id * 10 + int(ep[0])
                 ep = ep[1:]
         if ep:
-            self.data = json.loads(ep)
+            self.data = self.json.loads(ep)
         return attachment_count
 
     def reconstruct_binary(self, attachments):
diff --git a/socketio/server.py b/socketio/server.py
@@ -28,6 +28,10 @@ class Server(object):
                    ``bytes`` values are treated as binary.  This option has no
                    effect on Python 3, where text and binary payloads are
                    always automatically discovered.
+    :param json: An alternative json module to use for encoding and decoding
+                 packets. Custom json modules must have ``dumps`` and ``loads``
+                 functions that are compatible with the standard library
+                 versions.
     :param kwargs: Connection parameters for the underlying Engine.IO server.
 
     The Engine.IO configuration supports the following settings:
@@ -60,14 +64,17 @@ class Server(object):
                             ``False``.
     """
     def __init__(self, client_manager_class=None, logger=False, binary=False,
-                 **kwargs):
+                 json=None, **kwargs):
         if client_manager_class is None:
             client_manager_class = base_manager.BaseManager
         self.manager = client_manager_class(self)
         engineio_options = kwargs
         engineio_logger = engineio_options.pop('engineio_logger', None)
         if engineio_logger is not None:
             engineio_options['logger'] = engineio_logger
+        if json is not None:
+            packet.Packet.json = json
+            engineio_options['json'] = json
         self.eio = engineio.Server(**engineio_options)
         self.eio.on('connect', self._handle_eio_connect)
         self.eio.on('message', self._handle_eio_message)
diff --git a/tests/test_server.py b/tests/test_server.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import unittest
 
@@ -7,11 +8,16 @@
 else:
     import mock
 
+from socketio import packet
 from socketio import server
 
 
 @mock.patch('engineio.Server')
 class TestServer(unittest.TestCase):
+    def tearDown(self):
+        # restore JSON encoder, in case a test changed it
+        packet.Packet.json = json
+
     def test_create(self, eio):
         mgr = mock.MagicMock()
         s = server.Server(mgr, binary=True, foo='bar')
@@ -374,3 +380,28 @@ def test_logger(self, eio):
     def test_engineio_logger(self, eio):
         server.Server(engineio_logger='foo')
         eio.assert_called_once_with(**{'logger': 'foo'})
+
+    def test_custom_json(self, eio):
+        # Warning: this test cannot run in parallel with other tests, as it
+        # changes the JSON encoding/decoding functions
+
+        class CustomJSON(object):
+            @staticmethod
+            def dumps(*args, **kwargs):
+                return '*** encoded ***'
+
+            @staticmethod
+            def loads(*args, **kwargs):
+                return '+++ decoded +++'
+
+        server.Server(json=CustomJSON)
+        eio.assert_called_once_with(**{'json': CustomJSON})
+
+        pkt = packet.Packet(packet_type=packet.EVENT,
+                            data={six.text_type('foo'): six.text_type('bar')})
+        self.assertEqual(pkt.encode(), '2*** encoded ***')
+        pkt2 = packet.Packet(encoded_packet=pkt.encode())
+        self.assertEqual(pkt2.data, '+++ decoded +++')
+
+        # restore the default JSON module
+        packet.Packet.json = json
