Switch to asyncio.timeout if possible

Author: puddly

tests/api/test_request.py

@@ -1,8 +1,13 @@
+import sys
 import asyncio
 import logging
 
 import pytest
-import async_timeout
+
+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_znp.types as t
 import zigpy_znp.config as conf
@@ -66,7 +71,7 @@ async def test_cleanup_timeout_external(connected_znp):
 
     # This request will timeout because we didn't send anything back
     with pytest.raises(asyncio.TimeoutError):
-        async with async_timeout.timeout(0.1):
+        async with asyncio_timeout(0.1):
             await znp.request(c.UTIL.TimeAlive.Req())
 
     # We should be cleaned up
@@ -80,7 +85,7 @@ async def test_callback_rsp_cleanup_timeout_external(connected_znp):
 
     # This request will timeout because we didn't send anything back
     with pytest.raises(asyncio.TimeoutError):
-        async with async_timeout.timeout(0.1):
+        async with asyncio_timeout(0.1):
             await znp.request_callback_rsp(
                 request=c.UTIL.TimeAlive.Req(),
                 callback=c.SYS.ResetInd.Callback(partial=True),
@@ -262,7 +267,7 @@ async def test_znp_sreq_srsp(connected_znp):
 
     # Each SREQ must have a corresponding SRSP, so this will fail
     with pytest.raises(asyncio.TimeoutError):
-        async with async_timeout.timeout(0.5):
+        async with asyncio_timeout(0.5):
             await znp.request(c.SYS.Ping.Req())
 
     # This will work

tests/api/test_response.py

@@ -1,7 +1,12 @@
+import sys
 import asyncio
 
 import pytest
-import async_timeout
+
+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_znp.types as t
 import zigpy_znp.commands as c
@@ -61,7 +66,7 @@ async def send_soon(delay):
 
     asyncio.create_task(send_soon(0.1))
 
-    async with async_timeout.timeout(0.5):
+    async with asyncio_timeout(0.5):
         assert (await znp.wait_for_response(c.SYS.Ping.Rsp(partial=True))) == response
 
     # The response was successfully received so we should have no outstanding listeners
@@ -71,7 +76,7 @@ async def send_soon(delay):
     asyncio.create_task(send_soon(0.6))
 
     with pytest.raises(asyncio.TimeoutError):
-        async with async_timeout.timeout(0.5):
+        async with asyncio_timeout(0.5):
             assert (
                 await znp.wait_for_response(c.SYS.Ping.Rsp(partial=True))
             ) == response

zigpy_znp/api.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import sys
 import time
 import typing
 import asyncio
@@ -12,7 +13,12 @@
 from collections import Counter, defaultdict
 
 import zigpy.state
-import async_timeout
+
+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.zdo.types as zdo_t
 import zigpy.exceptions
 from zigpy.exceptions import NetworkNotFormed
@@ -296,7 +302,7 @@ async def start_network(self):
                         )
 
                 # Both versions still end with this callback
-                async with async_timeout.timeout(STARTUP_TIMEOUT):
+                async with asyncio_timeout(STARTUP_TIMEOUT):
                     await started_as_coordinator
             except asyncio.TimeoutError as e:
                 raise zigpy.exceptions.FormationFailure(
@@ -666,7 +672,7 @@ async def ping_task():
 
             # First, just try pinging
             try:
-                async with async_timeout.timeout(CONNECT_PING_TIMEOUT):
+                async with asyncio_timeout(CONNECT_PING_TIMEOUT):
                     return await self.request(c.SYS.Ping.Req())
             except asyncio.TimeoutError:
                 pass
@@ -682,7 +688,7 @@ async def ping_task():
             # At this point we have nothing left to try
             while True:
                 try:
-                    async with async_timeout.timeout(2 * CONNECT_PING_TIMEOUT):
+                    async with asyncio_timeout(2 * CONNECT_PING_TIMEOUT):
                         return await self.request(c.SYS.Ping.Req())
                 except asyncio.TimeoutError:
                     pass
@@ -691,7 +697,7 @@ async def ping_task():
             ping_task = asyncio.create_task(ping_task())
 
             try:
-                async with async_timeout.timeout(CONNECT_PROBE_TIMEOUT):
+                async with asyncio_timeout(CONNECT_PROBE_TIMEOUT):
                     result = await responses.get()
             except Exception:
                 ping_task.cancel()
@@ -706,7 +712,7 @@ async def ping_task():
                 LOGGER.debug("Giving ping task %0.2fs to finish", CONNECT_PING_TIMEOUT)
 
                 try:
-                    async with async_timeout.timeout(CONNECT_PING_TIMEOUT):
+                    async with asyncio_timeout(CONNECT_PING_TIMEOUT):
                         result = await ping_task  # type:ignore[misc]
                 except asyncio.TimeoutError:
                     ping_task.cancel()
@@ -1066,7 +1072,7 @@ async def request(
             self._uart.send(frame)
 
             # We should get a SRSP in a reasonable amount of time
-            async with async_timeout.timeout(
+            async with asyncio_timeout(
                 timeout or self._znp_config[conf.CONF_SREQ_TIMEOUT]
             ):
                 # We lock until either a sync response is seen or an error occurs
@@ -1108,7 +1114,7 @@ async def request_callback_rsp(
         # Typical request/response/callbacks are not backgrounded
         if not background:
             try:
-                async with async_timeout.timeout(timeout):
+                async with asyncio_timeout(timeout):
                     await self.request(request, timeout=timeout, **response_params)
 
                     return await callback_rsp
@@ -1119,7 +1125,7 @@ async def request_callback_rsp(
         start_time = time.monotonic()
 
         try:
-            async with async_timeout.timeout(timeout):
+            async with asyncio_timeout(timeout):
                 request_rsp = await self.request(request, **response_params)
         except Exception:
             # If the SREQ/SRSP pair fails, we must cancel the AREQ listener
@@ -1131,7 +1137,7 @@ async def request_callback_rsp(
         # the timeout
         async def callback_catcher(timeout):
             try:
-                async with async_timeout.timeout(timeout):
+                async with asyncio_timeout(timeout):
                     await callback_rsp
             finally:
                 self.remove_listener(listener)

zigpy_znp/tools/flash_read.py

@@ -2,7 +2,10 @@
 import asyncio
 import logging
 
-import async_timeout
+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_znp.commands as c
 from zigpy_znp.api import ZNP
@@ -14,7 +17,7 @@
 
 async def read_firmware(znp: ZNP) -> bytearray:
     try:
-        async with async_timeout.timeout(5):
+        async with asyncio_timeout(5):
             handshake_rsp = await znp.request_callback_rsp(
                 request=c.UBL.HandshakeReq.Req(),
                 callback=c.UBL.HandshakeRsp.Callback(partial=True),

zigpy_znp/tools/flash_write.py

@@ -4,7 +4,11 @@
 import asyncio
 import logging
 
-import async_timeout
+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_znp.types as t
 import zigpy_znp.commands as c
@@ -70,7 +74,7 @@ async def write_firmware(znp: ZNP, firmware: bytes, reset_nvram: bool):
         )
 
     try:
-        async with async_timeout.timeout(5):
+        async with asyncio_timeout(5):
             handshake_rsp = await znp.request_callback_rsp(
                 request=c.UBL.HandshakeReq.Req(),
                 callback=c.UBL.HandshakeRsp.Callback(partial=True),

zigpy_znp/zigbee/application.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import sys
 import asyncio
 import logging
 
@@ -11,7 +12,12 @@
 import zigpy.types
 import zigpy.config
 import zigpy.device
-import async_timeout
+
+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.profiles
 import zigpy.zdo.types as zdo_t
 import zigpy.application
@@ -64,7 +70,7 @@ def __init__(self, config: conf.ConfigType):
     @classmethod
     async def probe(cls, device_config):
         try:
-            async with async_timeout.timeout(PROBE_TIMEOUT):
+            async with asyncio_timeout(PROBE_TIMEOUT):
                 return await super().probe(device_config)
         except asyncio.TimeoutError:
             return False
@@ -609,7 +615,7 @@ async def _set_led_mode(self, *, led: t.uint8_t, mode: c.util.LEDMode) -> None:
 
         # XXX: If Z-Stack is not compiled with HAL_LED, it will just not respond at all
         try:
-            async with async_timeout.timeout(0.5):
+            async with asyncio_timeout(0.5):
                 await self._znp.request(
                     c.UTIL.LEDControl.Req(LED=led, Mode=mode),
                     RspStatus=t.Status.SUCCESS,
@@ -813,7 +819,7 @@ async def _send_request_raw(
             # Broadcasts and ZDO requests will not receive a confirmation
             await self._znp.request(request=request, RspStatus=t.Status.SUCCESS)
         else:
-            async with async_timeout.timeout(
+            async with asyncio_timeout(
                 EXTENDED_DATA_CONFIRM_TIMEOUT
                 if extended_timeout
                 else DATA_CONFIRM_TIMEOUT