Replace ujson with orjson

examples/langserver.py

@@ -5,8 +5,8 @@
 from pylsp_jsonrpc import dispatchers, endpoint
 
 try:
-    import ujson as json
-except Exception:  # pylint: disable=broad-except
+    import orjson as json
+except ImportError:
     import json
 
 log = logging.getLogger(__name__)

examples/langserver_ext.py

@@ -7,8 +7,8 @@
 from pylsp_jsonrpc import streams
 
 try:
-    import ujson as json
-except Exception:  # pylint: disable=broad-except
+    import orjson as json
+except ImportError:
     import json
 
 log = logging.getLogger(__name__)

pylsp_jsonrpc/streams.py

@@ -3,10 +3,11 @@
 
 import logging
 import threading
+import sys
 
 try:
-    import ujson as json
-except Exception:  # pylint: disable=broad-except
+    import orjson as json
+except ImportError:
     import json
 
 log = logging.getLogger(__name__)
@@ -83,7 +84,16 @@ class JsonRpcStreamWriter:
     def __init__(self, wfile, **json_dumps_args):
         self._wfile = wfile
         self._wfile_lock = threading.Lock()
-        self._json_dumps_args = json_dumps_args
+
+        if 'orjson' in sys.modules and json_dumps_args.pop('sort_keys'):
+            # orjson needs different option handling;
+            # pylint has an erroneous error here https://github.com/pylint-dev/pylint/issues/9762
+            self._json_dumps_args = {'option': json.OPT_SORT_KEYS}  # pylint: disable=maybe-no-member
+            self._json_dumps_args.update(**json_dumps_args)
+        else:
+            self._json_dumps_args = json_dumps_args
+            # omit unnecessary whitespace for consistency with orjson
+            self._json_dumps_args.setdefault('separators', (',', ':'))
 
     def close(self):
         with self._wfile_lock:
@@ -96,16 +106,16 @@ def write(self, message):
             try:
                 body = json.dumps(message, **self._json_dumps_args)
 
-                # Ensure we get the byte length, not the character length
-                content_length = len(body) if isinstance(body, bytes) else len(body.encode('utf-8'))
+                # orjson gives bytes, builtin json gives str. ensure we have bytes
+                body_bytes = body if isinstance(body, bytes) else body.encode('utf-8')
 
                 response = (
-                    f"Content-Length: {content_length}\r\n"
-                    f"Content-Type: application/vscode-jsonrpc; charset=utf8\r\n\r\n"
-                    f"{body}"
-                )
+                    b"Content-Length: %(length)i\r\n"
+                    b"Content-Type: application/vscode-jsonrpc; charset=utf8\r\n\r\n"
+                    b"%(body)s"
+                ) % {b'length': len(body_bytes), b'body': body_bytes}
 
-                self._wfile.write(response.encode('utf-8'))
+                self._wfile.write(response)
                 self._wfile.flush()
             except Exception:  # pylint: disable=broad-except
                 log.exception("Failed to write message to output file %s", message)

pyproject.toml

@@ -11,7 +11,7 @@ authors = [{name = "Python Language Server Contributors"}]
 description = "JSON RPC 2.0 server library"
 license = {text = "MIT"}
 requires-python = ">=3.8"
-dependencies = ["ujson>=3.0.0"]
+dependencies = ["orjson>=3.10.0"]
 dynamic = ["version"]
 classifiers = [
     "License :: OSI Approved :: MIT License",

test/test_streams.py

@@ -7,6 +7,7 @@
 import datetime
 import sys
 from unittest import mock
+import json
 import pytest
 
 from pylsp_jsonrpc.streams import JsonRpcStreamReader, JsonRpcStreamWriter
@@ -77,25 +78,48 @@ def test_reader_bad_json(rfile, reader):
 
 
 def test_writer(wfile, writer):
-    writer.write({
+    data = {
         'id': 'hello',
         'method': 'method',
         'params': {}
-    })
-    if 'ujson' in sys.modules:
-        assert wfile.getvalue() == (
-            b'Content-Length: 44\r\n'
-            b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
-            b'\r\n'
-            b'{"id":"hello","method":"method","params":{}}'
-        )
-    else:
-        assert wfile.getvalue() == (
-            b'Content-Length: 49\r\n'
-            b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
-            b'\r\n'
-            b'{"id": "hello", "method": "method", "params": {}}'
-        )
+    }
+    writer.write(data)
+
+    assert wfile.getvalue() == (
+        b'Content-Length: 44\r\n'
+        b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
+        b'\r\n'
+        b'{"id":"hello","method":"method","params":{}}'
+    )
+
+
+def test_writer_stdlib_json(wfile):
+    """Test the stream writer using the standard json lib."""
+    data = {
+        'id': 'hello',
+        'method': 'method',
+        'params': {}
+    }
+    orig_modules = sys.modules
+    try:
+        # Pretend orjson wasn't imported when initializing the writer.
+        sys.modules = {'json': json}
+        std_json_writer = JsonRpcStreamWriter(wfile, sort_keys=True)
+    finally:
+        sys.modules = orig_modules
+
+    with mock.patch('pylsp_jsonrpc.streams.json') as streams_json:
+        # Mock the imported json's dumps function to use the stdlib's dumps,
+        # whether orjson is available or not.
+        streams_json.dumps = json.dumps
+        std_json_writer.write(data)
+
+    assert wfile.getvalue() == (
+        b'Content-Length: 44\r\n'
+        b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
+        b'\r\n'
+        b'{"id":"hello","method":"method","params":{}}'
+    )
 
 
 class JsonDatetime(datetime.datetime):