Switch asyncio.wait_for to asyncio_timeout

see zigpy/zigpy#1187
zigpy · Sep 4, 2023 · 9a19262 · 9a19262
1 parent 09e952b
commit 9a19262
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/zigpy_deconz/api.py b/zigpy_deconz/api.py
@@ -9,6 +9,7 @@
 import logging
 from typing import Any, Callable
 
+from zigpy.application import asyncio_timeout
 from zigpy.config import CONF_DEVICE_PATH
 import zigpy.exceptions
 from zigpy.types import APSStatus, Bool, Channels
@@ -303,7 +304,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 +392,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
@@ -9,6 +9,7 @@
 from typing import Any
 
 import zigpy.application
+from zigpy.application import asyncio_timeout
 import zigpy.config
 import zigpy.device
 import zigpy.endpoint
@@ -140,7 +141,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 +452,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 +553,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)