diff --git a/examples/langserver.py b/examples/langserver.py index a1ddc8f..5a40f1c 100644 --- a/examples/langserver.py +++ b/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__) diff --git a/examples/langserver_ext.py b/examples/langserver_ext.py index 3371a20..19d227c 100644 --- a/examples/langserver_ext.py +++ b/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__) diff --git a/pylsp_jsonrpc/streams.py b/pylsp_jsonrpc/streams.py index 40048a9..7dacbea 100644 --- a/pylsp_jsonrpc/streams.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 376785c..2c4d196 100644 --- a/pyproject.toml +++ b/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", diff --git a/test/test_streams.py b/test/test_streams.py index 14fe1bb..aff10da 100644 --- a/test/test_streams.py +++ b/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):