Migrate http client to use aiohttp instead of httpx (#643)

kasa/httpclient.py

@@ -1,15 +1,16 @@
 """Module for HttpClientSession class."""
-import logging
-from typing import Any, Dict, Optional, Tuple, Type, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
-import httpx
+import aiohttp
 
 from .deviceconfig import DeviceConfig
 from .exceptions import ConnectionException, SmartDeviceException, TimeoutException
+from .json import loads as json_loads
 
-logging.getLogger("httpx").propagate = False
 
-InnerHttpType = Type[httpx.AsyncClient]
+def get_cookie_jar() -> aiohttp.CookieJar:
+    """Return a new cookie jar with the correct options for device communication."""
+    return aiohttp.CookieJar(unsafe=True, quote_cookie=False)
 
 
 class HttpClient:
@@ -17,18 +18,20 @@ class HttpClient:
 
     def __init__(self, config: DeviceConfig) -> None:
         self._config = config
-        self._client: httpx.AsyncClient = None
+        self._client: aiohttp.ClientSession = None
+        self._jar = aiohttp.CookieJar(unsafe=True, quote_cookie=False)
+        self._last_url = f"http://{self._config.host}/"
 
     @property
-    def client(self) -> httpx.AsyncClient:
+    def client(self) -> aiohttp.ClientSession:
         """Return the underlying http client."""
         if self._config.http_client and issubclass(
-            self._config.http_client.__class__, httpx.AsyncClient
+            self._config.http_client.__class__, aiohttp.ClientSession
         ):
             return self._config.http_client
 
         if not self._client:
-            self._client = httpx.AsyncClient()
+            self._client = aiohttp.ClientSession(cookie_jar=get_cookie_jar())
         return self._client
 
     async def post(
@@ -43,12 +46,8 @@ class HttpClient:
     ) -> Tuple[int, Optional[Union[Dict, bytes]]]:
         """Send an http post request to the device."""
         response_data = None
-        cookies = None
-        if cookies_dict:
-            cookies = httpx.Cookies()
-            for name, value in cookies_dict.items():
-                cookies.set(name, value)
-        self.client.cookies.clear()
+        self._last_url = url
+        self.client.cookie_jar.clear()
         try:
             resp = await self.client.post(
                 url,
@@ -56,14 +55,14 @@ class HttpClient:
                 data=data,
                 json=json,
                 timeout=self._config.timeout,
-                cookies=cookies,
+                cookies=cookies_dict,
                 headers=headers,
             )
-        except httpx.ConnectError as ex:
+        except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as ex:
             raise ConnectionException(
                 f"Unable to connect to the device: {self._config.host}: {ex}"
             ) from ex
-        except httpx.TimeoutException as ex:
+        except aiohttp.ServerTimeoutError as ex:
             raise TimeoutException(
                 "Unable to query the device, " + f"timed out: {self._config.host}: {ex}"
             ) from ex
@@ -72,18 +71,25 @@ class HttpClient:
                 f"Unable to query the device: {self._config.host}: {ex}"
             ) from ex
 
-        if resp.status_code == 200:
-            response_data = resp.json() if json else resp.content
+        async with resp:
+            if resp.status == 200:
+                response_data = await resp.read()
+                if json:
+                    response_data = json_loads(response_data.decode())
 
-        return resp.status_code, response_data
+        return resp.status, response_data
 
-    def get_cookie(self, cookie_name: str) -> str:
+    def get_cookie(self, cookie_name: str) -> Optional[str]:
         """Return the cookie with cookie_name."""
-        return self._client.cookies.get(cookie_name)
+        if cookie := self.client.cookie_jar.filter_cookies(self._last_url).get(
+            cookie_name
+        ):
+            return cookie.value
+        return None
 
     async def close(self) -> None:
-        """Close the protocol."""
+        """Close the client."""
         client = self._client
         self._client = None
         if client:
-            await client.aclose()
+            await client.close()