diff --git a/kasa/protocol.py b/kasa/protocol.py index cd9066c6f..461dd85ad 100755 --- a/kasa/protocol.py +++ b/kasa/protocol.py @@ -17,6 +17,10 @@ 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 @@ async def _connect(self, timeout: int) -> None: 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 @@ async def _query(self, request: str, retry_count: int, timeout: int) -> Dict: 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 c846dc1a8..23e4b3638 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 68a2159d3..c905cf915 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 = "*"