Use orjson when already installed or with speedups extra (#466)

* Use orjson when already installed

* Use orjson when already installed

* fix patch target

* fix patch target

* add speedups extra

* Update README.md

* Update README.md

README.md

@@ -18,6 +18,11 @@ You can install the most recent release using pip:
 pip install python-kasa
 ```
 
+If you are using cpython, it is recommended to install with `[speedups]` to enable orjson (faster json support):
+```
+pip install python-kasa[speedups]
+```
+
 Alternatively, you can clone this repository and use poetry to install the development version:
 ```
 git clone https://github.com/python-kasa/python-kasa.git

kasa/discover.py

@@ -1,10 +1,11 @@
 """Discovery module for TP-Link Smart Home devices."""
 import asyncio
-import json
 import logging
 import socket
 from typing import Awaitable, Callable, Dict, Optional, Type, cast
 
+from kasa.json import dumps as json_dumps
+from kasa.json import loads as json_loads
 from kasa.protocol import TPLinkSmartHomeProtocol
 from kasa.smartbulb import SmartBulb
 from kasa.smartdevice import SmartDevice, SmartDeviceException
@@ -63,7 +64,7 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
 
     def do_discover(self) -> None:
         """Send number of discovery datagrams."""
-        req = json.dumps(Discover.DISCOVERY_QUERY)
+        req = json_dumps(Discover.DISCOVERY_QUERY)
         _LOGGER.debug("[DISCOVERY] %s >> %s", self.target, Discover.DISCOVERY_QUERY)
         encrypted_req = TPLinkSmartHomeProtocol.encrypt(req)
         for i in range(self.discovery_packets):
@@ -75,7 +76,7 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
         if ip in self.discovered_devices:
             return
 
-        info = json.loads(TPLinkSmartHomeProtocol.decrypt(data))
+        info = json_loads(TPLinkSmartHomeProtocol.decrypt(data))
         _LOGGER.debug("[DISCOVERY] %s << %s", ip, info)
 
         try:

kasa/json.py

@@ -0,0 +1,15 @@
+"""JSON abstraction."""
+
+try:
+    import orjson
+
+    def dumps(obj, *, default=None):
+        """Dump JSON."""
+        return orjson.dumps(obj).decode()
+
+    loads = orjson.loads
+except ImportError:
+    import json
+
+    dumps = json.dumps
+    loads = json.loads

kasa/protocol.py

@@ -12,13 +12,14 @@ http://www.apache.org/licenses/LICENSE-2.0
 import asyncio
 import contextlib
 import errno
-import json
 import logging
 import struct
 from pprint import pformat as pf
 from typing import Dict, Generator, Optional, Union
 
 from .exceptions import SmartDeviceException
+from .json import dumps as json_dumps
+from .json import loads as json_loads
 
 _LOGGER = logging.getLogger(__name__)
 _NO_RETRY_ERRORS = {errno.EHOSTDOWN, errno.EHOSTUNREACH, errno.ECONNREFUSED}
@@ -64,7 +65,7 @@ class TPLinkSmartHomeProtocol:
             self.query_lock = asyncio.Lock()
 
         if isinstance(request, dict):
-            request = json.dumps(request)
+            request = json_dumps(request)
             assert isinstance(request, str)
 
         timeout = TPLinkSmartHomeProtocol.DEFAULT_TIMEOUT
@@ -96,7 +97,7 @@ class TPLinkSmartHomeProtocol:
 
         buffer = await self.reader.readexactly(length)
         response = TPLinkSmartHomeProtocol.decrypt(buffer)
-        json_payload = json.loads(response)
+        json_payload = json_loads(response)
         if debug_log:
             _LOGGER.debug("%s << %s", self.host, pf(json_payload))
 

kasa/tests/test_discovery.py

@@ -91,7 +91,7 @@ async def test_discover_send(mocker):
 async def test_discover_datagram_received(mocker, discovery_data):
     """Verify that datagram received fills discovered_devices."""
     proto = _DiscoverProtocol()
-    mocker.patch("json.loads", return_value=discovery_data)
+    mocker.patch("kasa.discover.json_loads", return_value=discovery_data)
     mocker.patch.object(protocol.TPLinkSmartHomeProtocol, "encrypt")
     mocker.patch.object(protocol.TPLinkSmartHomeProtocol, "decrypt")
 
@@ -109,7 +109,7 @@ async def test_discover_datagram_received(mocker, discovery_data):
 async def test_discover_invalid_responses(msg, data, mocker):
     """Verify that we don't crash whole discovery if some devices in the network are sending unexpected data."""
     proto = _DiscoverProtocol()
-    mocker.patch("json.loads", return_value=data)
+    mocker.patch("kasa.discover.json_loads", return_value=data)
     mocker.patch.object(protocol.TPLinkSmartHomeProtocol, "encrypt")
     mocker.patch.object(protocol.TPLinkSmartHomeProtocol, "decrypt")
 

pyproject.toml

@@ -24,6 +24,9 @@ anyio = "*"  # see https://github.com/python-trio/asyncclick/issues/18
 asyncclick = ">=8"
 pydantic = "^1"
 
+# speed ups
+orjson = { "version" = ">=3.9.1", optional = true, extras = ["speedups"] }
+
 # required only for docs
 sphinx = { version = "^4", optional = true }
 sphinx_rtd_theme = { version = "^0", optional = true }