From 117a7ac64a2b7d419692970fba09f27c3a4290c2 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Fri, 21 Jul 2023 04:50:54 -0500
Subject: [PATCH] Replace asyncio.wait_for with async-timeout (#480)

asyncio.wait_for has some underlying problems that are only fixed in cpython 3.12.

Use async_timeout instead until the minimum supported version is 3.11+ and it can be replaced with asyncio.timeout

See https://github.com/python/cpython/pull/98518
---
 kasa/protocol.py | 13 +++++++++----
 poetry.lock      | 14 +++++++++++++-
 pyproject.toml   |  1 +
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/kasa/protocol.py b/kasa/protocol.py
index cd9066c6..461dd85a 100755
--- a/kasa/protocol.py
+++ b/kasa/protocol.py
@@ -17,6 +17,10 @@ import struct
 from pprint import pformat as pf
 from typing import Dict, Generator, Optional, Union
 
+# When support for cpython older than 3.11 is dropped
+# async_timeout can be replaced with asyncio.timeout
+from async_timeout import timeout as asyncio_timeout
+
 from .exceptions import SmartDeviceException
 from .json import dumps as json_dumps
 from .json import loads as json_loads
@@ -79,8 +83,10 @@ class TPLinkSmartHomeProtocol:
         if self.writer:
             return
         self.reader = self.writer = None
+
         task = asyncio.open_connection(self.host, self.port)
-        self.reader, self.writer = await asyncio.wait_for(task, timeout=timeout)
+        async with asyncio_timeout(timeout):
+            self.reader, self.writer = await task
 
     async def _execute_query(self, request: str) -> Dict:
         """Execute a query on the device and wait for the response."""
@@ -155,9 +161,8 @@ class TPLinkSmartHomeProtocol:
             try:
                 assert self.reader is not None
                 assert self.writer is not None
-                return await asyncio.wait_for(
-                    self._execute_query(request), timeout=timeout
-                )
+                async with asyncio_timeout(timeout):
+                    return await self._execute_query(request)
             except Exception as ex:
                 await self.close()
                 if retry >= retry_count:
diff --git a/poetry.lock b/poetry.lock
index c846dc1a..23e4b363 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -34,6 +34,18 @@ doc = ["Sphinx (>=6.1.0)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "s
 test = ["anyio[trio]", "coverage[toml] (>=4.5)", "hypothesis (>=4.0)", "mock (>=4)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"]
 trio = ["trio (<0.22)"]
 
+[[package]]
+name = "async-timeout"
+version = "4.0.2"
+description = "Timeout context manager for asyncio programs"
+category = "main"
+optional = false
+python-versions = ">=3.6"
+files = [
+    {file = "async-timeout-4.0.2.tar.gz", hash = "sha256:2163e1640ddb52b7a8c80d0a67a08587e5d245cc9c553a74a847056bc2976b15"},
+    {file = "async_timeout-4.0.2-py3-none-any.whl", hash = "sha256:8ca1e4fcf50d07413d66d1a5e416e42cfdf5851c981d679a09851a6853383b3c"},
+]
+
 [[package]]
 name = "asyncclick"
 version = "8.1.3.4"
@@ -1404,4 +1416,4 @@ speedups = ["orjson", "kasa-crypt"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8"
-content-hash = "b85f55f0ca928b1f3510da37196c21f40eb07cd4d07b3a9c3dd29215ba9777fe"
+content-hash = "fcb657fbabe28548021f5f6a1fcb7b60aa82d60f3de015b4b0c7b37260a6a29f"
diff --git a/pyproject.toml b/pyproject.toml
index 68a2159d..c905cf91 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,6 +34,7 @@ sphinx_rtd_theme = { version = "^0", optional = true }
 sphinxcontrib-programoutput = { version = "^0", optional = true }
 myst-parser = { version = "*", optional = true }
 docutils = { version = ">=0.17", optional = true }
+async-timeout = ">=3.0.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "*"