From 88d825a803917ab3e0e9797df9f47be7ed79bb9c Mon Sep 17 00:00:00 2001 From: Rapptz Date: Mon, 5 Jul 2021 04:01:19 -0400 Subject: [PATCH] Allow use of orjson instead of json The difference in speed seems negligible at start up, which is when most time is taken for actually parsing JSON. I could potentially be missing something but profiling didn't point to any discernable difference. --- discord/gateway.py | 5 ++--- discord/http.py | 2 +- discord/utils.py | 23 +++++++++++++++++++++-- setup.py | 3 +++ 4 files changed, 27 insertions(+), 6 deletions(-) diff --git a/discord/gateway.py b/discord/gateway.py index 5fec87484..0a6b05d05 100644 --- a/discord/gateway.py +++ b/discord/gateway.py @@ -25,7 +25,6 @@ DEALINGS IN THE SOFTWARE. import asyncio from collections import namedtuple, deque import concurrent.futures -import json import logging import struct import sys @@ -421,7 +420,7 @@ class DiscordWebSocket: msg = self._zlib.decompress(self._buffer) msg = msg.decode('utf-8') self._buffer = bytearray() - msg = json.loads(msg) + msg = utils.from_json(msg) log.debug('For Shard ID %s: WebSocket Event: %s', self.shard_id, msg) self._dispatch('socket_response', msg) @@ -882,7 +881,7 @@ class DiscordVoiceWebSocket: # This exception is handled up the chain msg = await asyncio.wait_for(self.ws.receive(), timeout=30.0) if msg.type is aiohttp.WSMsgType.TEXT: - await self.received_message(json.loads(msg.data)) + await self.received_message(utils.from_json(msg.data)) elif msg.type is aiohttp.WSMsgType.ERROR: log.debug('Received %s', msg) raise ConnectionClosed(self.ws, shard_id=None) from msg.data diff --git a/discord/http.py b/discord/http.py index cfc71e22c..c9c340aea 100644 --- a/discord/http.py +++ b/discord/http.py @@ -99,7 +99,7 @@ async def json_or_text(response: aiohttp.ClientResponse) -> Union[Dict[str, Any] text = await response.text(encoding='utf-8') try: if response.headers['content-type'] == 'application/json': - return json.loads(text) + return utils.from_json(text) except KeyError: # Thanks Cloudflare pass diff --git a/discord/utils.py b/discord/utils.py index 6070882f2..b5563a672 100644 --- a/discord/utils.py +++ b/discord/utils.py @@ -63,6 +63,14 @@ import warnings from .errors import InvalidArgument +try: + import orjson +except ModuleNotFoundError: + HAS_ORJSON = False +else: + HAS_ORJSON = True + + __all__ = ( 'oauth_url', 'snowflake_time', @@ -468,8 +476,19 @@ def _bytes_to_base64_data(data: bytes) -> str: return fmt.format(mime=mime, data=b64) -def to_json(obj: Any) -> str: - return json.dumps(obj, separators=(',', ':'), ensure_ascii=True) +if HAS_ORJSON: + + def to_json(obj: Any) -> str: # type: ignore + return orjson.dumps(obj).decode('utf-8') + + from_json = orjson.loads # type: ignore + +else: + + def to_json(obj: Any) -> str: + return json.dumps(obj, separators=(',', ':'), ensure_ascii=True) + + from_json = json.loads def _parse_ratelimit_header(request: Any, *, use_clock: bool = False) -> float: diff --git a/setup.py b/setup.py index bdf1549f5..251241b4a 100644 --- a/setup.py +++ b/setup.py @@ -39,6 +39,9 @@ extras_require = { 'sphinx==4.0.2', 'sphinxcontrib_trio==1.1.2', 'sphinxcontrib-websupport', + ], + 'speed': [ + 'orjson>=3.5.4', ] }