Switch asyncio.wait_for to asyncio_timeout (#228)

* Switch asyncio.wait_for to asyncio_timeout see zigpy/zigpy#1187 * Directly import from the `async_timeout` package --------- Co-authored-by: puddly <32534428+puddly@users.noreply.github.com>
zigpy · Sep 11, 2023 · dc4a1d3 · dc4a1d3
1 parent 09e952b
commit dc4a1d3
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/zigpy_deconz/api.py b/zigpy_deconz/api.py
@@ -7,8 +7,14 @@
 import enum
 import functools
 import logging
+import sys
 from typing import Any, Callable
 
+if sys.version_info[:2] < (3, 11):
+    from async_timeout import timeout as asyncio_timeout  # pragma: no cover
+else:
+    from asyncio import timeout as asyncio_timeout  # pragma: no cover
+
 from zigpy.config import CONF_DEVICE_PATH
 import zigpy.exceptions
 from zigpy.types import APSStatus, Bool, Channels
@@ -303,7 +309,8 @@ async def _command(self, cmd, *args):
             fut = asyncio.Future()
             self._awaiting[seq] = fut
             try:
-                return await asyncio.wait_for(fut, timeout=COMMAND_TIMEOUT)
+                async with asyncio_timeout(COMMAND_TIMEOUT):
+                    return await fut
             except asyncio.TimeoutError:
                 LOGGER.warning(
                     "No response to '%s' command with seq id '0x%02x'", cmd, seq
@@ -390,7 +397,8 @@ async def probe(cls, device_config: dict[str, Any]) -> bool:
         """Probe port for the device presence."""
         api = cls(None, device_config)
         try:
-            await asyncio.wait_for(api._probe(), timeout=PROBE_TIMEOUT)
+            async with asyncio_timeout(PROBE_TIMEOUT):
+                await api._probe()
             return True
         except Exception as exc:
             LOGGER.debug(

diff --git a/zigpy_deconz/zigbee/application.py b/zigpy_deconz/zigbee/application.py
@@ -6,8 +6,14 @@
 import importlib.metadata
 import logging
 import re
+import sys
 from typing import Any
 
+if sys.version_info[:2] < (3, 11):
+    from async_timeout import timeout as asyncio_timeout  # pragma: no cover
+else:
+    from asyncio import timeout as asyncio_timeout  # pragma: no cover
+
 import zigpy.application
 import zigpy.config
 import zigpy.device
@@ -140,7 +146,8 @@ async def change_loop():
         await self._api.change_network_state(target_state)
 
         try:
-            await asyncio.wait_for(change_loop(), timeout=timeout)
+            async with asyncio_timeout(timeout):
+                await change_loop()
         except asyncio.TimeoutError:
             if target_state != NetworkState.CONNECTED:
                 raise
@@ -450,7 +457,8 @@ async def send_packet(self, packet):
                         f"Failed to enqueue packet: {ex!r}", ex.status
                     )
 
-                status = await asyncio.wait_for(req.result, SEND_CONFIRM_TIMEOUT)
+                async with asyncio_timeout(SEND_CONFIRM_TIMEOUT):
+                    status = await req.result
 
                 if status != TXStatus.SUCCESS:
                     raise zigpy.exceptions.DeliveryError(
@@ -550,8 +558,10 @@ async def _reconnect_loop(self) -> None:
             LOGGER.debug("Reconnecting, attempt %s", attempt)
 
             try:
-                await asyncio.wait_for(self.connect(), timeout=10)
-                await asyncio.wait_for(self.initialize(), timeout=10)
+                async with asyncio_timeout(10):
+                    await self.connect()
+                async with asyncio_timeout(10):
+                    await self.initialize()
                 break
             except Exception as exc:
                 wait = 2 ** min(attempt, 5)