Add an API to choose the underlying JSON encoder.

.gitignore

@@ -6,3 +6,4 @@ __pycache__/
 *.pyc
 /build/
 /dist/
+.python-version

canonicaljson.py

@@ -20,17 +20,6 @@
 
 from frozendict import frozendict
 
-if platform.python_implementation() == "PyPy":  # pragma: no cover
-    # pypy ships with an optimised JSON encoder/decoder that is faster than
-    # simplejson's C extension.
-    import json
-else:  # pragma: no cover
-    # using simplejson rather than regular json on CPython gives approximately
-    # a 100% performance improvement (as measured on python 2.7.12/simplejson
-    # 3.13.2)
-    import simplejson as json
-
-
 __version__ = '1.1.4'
 
 
@@ -50,19 +39,37 @@ def _default(obj):
 # (in any case, simplejson's ensure_ascii doesn't get U+2028 and U+2029 right,
 # as per https://github.com/simplejson/simplejson/issues/206).
 #
-_canonical_encoder = json.JSONEncoder(
-    ensure_ascii=True,
-    separators=(',', ':'),
-    sort_keys=True,
-    default=_default,
-)
-
-_pretty_encoder = json.JSONEncoder(
-    ensure_ascii=True,
-    indent=4,
-    sort_keys=True,
-    default=_default,
-)
+
+# Declare these in the module scope, but they get configured in
+# set_json_library.
+_canonical_encoder = None
+_pretty_encoder = None
+
+
+def set_json_library(json_lib):
+    """
+    Set the underlying JSON library that canonicaljson uses to json_lib.
+
+    Params:
+        json_lib: The module to use for JSON encoding, Must have a
+            `JSONEncoder` property.
+    """
+    global _canonical_encoder
+    _canonical_encoder = json_lib.JSONEncoder(
+        ensure_ascii=True,
+        separators=(',', ':'),
+        sort_keys=True,
+        default=_default,
+    )
+
+    global _pretty_encoder
+    _pretty_encoder = json_lib.JSONEncoder(
+        ensure_ascii=True,
+        indent=4,
+        sort_keys=True,
+        default=_default,
+    )
+
 
 # This regexp matches either `\uNNNN` or `\\`. We match '\\' (and leave it
 # unchanged) to make sure that the regex doesn't accidentally capture the uNNNN
@@ -157,3 +164,20 @@ def encode_pretty_printed_json(json_object):
     """Encodes the JSON object dict as human readable ascii bytes."""
 
     return _pretty_encoder.encode(json_object).encode("ascii")
+
+
+if platform.python_implementation() == "PyPy":  # pragma: no cover
+    # pypy ships with an optimised JSON encoder/decoder that is faster than
+    # simplejson's C extension.
+    import json
+else:  # pragma: no cover
+    # using simplejson rather than regular json on CPython for backwards
+    # compatibility (simplejson on Python 3.5 handles parsing of bytes while
+    # the standard library json does not).
+    #
+    # Note that it seems performance is on par or better using json from the
+    # standard library as of Python 3.7.
+    import simplejson as json
+
+# Set the JSON library to the backwards compatible version.
+set_json_library(json)

test_canonicaljson.py

@@ -18,11 +18,13 @@
 from canonicaljson import (
     encode_canonical_json,
     encode_pretty_printed_json,
+    set_json_library,
 )
 
 from frozendict import frozendict
 
 import unittest
+from unittest import mock
 
 
 class TestCanonicalJson(unittest.TestCase):
@@ -78,3 +80,15 @@ class Unknown(object):
         unknown_object = Unknown()
         with self.assertRaises(Exception):
             encode_canonical_json(unknown_object)
+
+    def test_set_json(self):
+        """Ensure that changing the underlying JSON implementation works."""
+        mock_json = mock.Mock(spec=["JSONEncoder"])
+        mock_json.JSONEncoder.return_value.encode.return_value = "sentinel"
+        try:
+            set_json_library(mock_json)
+            self.assertEqual(encode_canonical_json({}), b"sentinel")
+        finally:
+            # Reset the JSON library to whatever was originally set.
+            from canonicaljson import json
+            set_json_library(json)