|
14 | 14 | import errno
|
15 | 15 | import logging
|
16 | 16 | import struct
|
| 17 | +import sys |
17 | 18 | from pprint import pformat as pf
|
18 | 19 | from typing import Dict, Generator, Optional, Union
|
19 | 20 |
|
20 | 21 | from .exceptions import SmartDeviceException
|
21 | 22 | from .json import dumps as json_dumps
|
22 | 23 | from .json import loads as json_loads
|
23 | 24 |
|
| 25 | +if sys.version_info[:2] < (3, 11): |
| 26 | + from async_timeout import timeout as asyncio_timeout |
| 27 | +else: |
| 28 | + from asyncio import timeout as asyncio_timeout |
| 29 | + |
| 30 | + |
24 | 31 | _LOGGER = logging.getLogger(__name__)
|
25 | 32 | _NO_RETRY_ERRORS = {errno.EHOSTDOWN, errno.EHOSTUNREACH, errno.ECONNREFUSED}
|
26 | 33 |
|
@@ -79,8 +86,10 @@ async def _connect(self, timeout: int) -> None:
|
79 | 86 | if self.writer:
|
80 | 87 | return
|
81 | 88 | self.reader = self.writer = None
|
| 89 | + |
82 | 90 | task = asyncio.open_connection(self.host, self.port)
|
83 |
| - self.reader, self.writer = await asyncio.wait_for(task, timeout=timeout) |
| 91 | + async with asyncio_timeout(timeout): |
| 92 | + self.reader, self.writer = await task |
84 | 93 |
|
85 | 94 | async def _execute_query(self, request: str) -> Dict:
|
86 | 95 | """Execute a query on the device and wait for the response."""
|
@@ -155,9 +164,8 @@ async def _query(self, request: str, retry_count: int, timeout: int) -> Dict:
|
155 | 164 | try:
|
156 | 165 | assert self.reader is not None
|
157 | 166 | assert self.writer is not None
|
158 |
| - return await asyncio.wait_for( |
159 |
| - self._execute_query(request), timeout=timeout |
160 |
| - ) |
| 167 | + async with asyncio_timeout(timeout): |
| 168 | + await self._execute_query(request) |
161 | 169 | except Exception as ex:
|
162 | 170 | await self.close()
|
163 | 171 | if retry >= retry_count:
|
|
0 commit comments