Switch from asyncio.wait_for to async_timeout to avoid task creation (#…

…387)

flux_led/aiodevice.py

@@ -6,6 +6,7 @@
 
 from .aioprotocol import AIOLEDENETProtocol
 from .aioscanner import AIOBulbScanner
+from .aioutils import asyncio_timeout
 from .base_device import (
     ALL_ADDRESSABLE_PROTOCOLS,
     ALL_IC_PROTOCOLS,
@@ -133,7 +134,8 @@ async def _async_remote_config_setup(self) -> None:
         assert self._protocol is not None
         await self._async_send_msg(self._protocol.construct_query_remote_config())
         try:
-            await asyncio.wait_for(self._remote_config_future, timeout=self.timeout)
+            async with asyncio_timeout(self.timeout):
+                await self._remote_config_future
         except asyncio.TimeoutError:
             _LOGGER.warning("%s: Could not determine 2.4ghz remote config", self.ipaddr)
 
@@ -142,7 +144,8 @@ async def _async_switch_setup(self) -> None:
         assert self._protocol is not None
         await self._async_send_msg(self._protocol.construct_power_restore_state_query())
         try:
-            await asyncio.wait_for(self._power_restore_future, timeout=self.timeout)
+            async with asyncio_timeout(self.timeout):
+                await self._power_restore_future
         except asyncio.TimeoutError:
             self.set_unavailable()
             raise RuntimeError(
@@ -163,7 +166,8 @@ async def _async_device_config_setup(self) -> None:
         assert isinstance(self._protocol, ALL_ADDRESSABLE_PROTOCOLS)
         await self._async_send_msg(self._protocol.construct_request_strip_setting())
         try:
-            await asyncio.wait_for(self._device_config_future, timeout=self.timeout)
+            async with asyncio_timeout(self.timeout):
+                await self._device_config_future
         except asyncio.TimeoutError:
             self.set_unavailable()
             raise RuntimeError(f"{self.ipaddr}: Could not determine number pixels")
@@ -599,7 +603,8 @@ async def async_get_time(self) -> Optional[datetime]:
         async with self._get_time_lock:
             self._get_time_future = asyncio.Future()
             try:
-                await asyncio.wait_for(self._get_time_future, timeout=self.timeout)
+                async with asyncio_timeout(self.timeout):
+                    await self._get_time_future
             except asyncio.TimeoutError:
                 _LOGGER.warning("%s: Could not get time from the device", self.ipaddr)
                 return None
@@ -615,7 +620,8 @@ async def async_get_timers(self) -> Optional[List[LedTimer]]:
         async with self._get_timers_lock:
             self._get_timers_future = asyncio.Future()
             try:
-                await asyncio.wait_for(self._get_timers_future, timeout=self.timeout)
+                async with asyncio_timeout(self.timeout):
+                    await self._get_timers_future
             except asyncio.TimeoutError:
                 _LOGGER.warning("%s: Could not get timers from the device", self.ipaddr)
                 return None
@@ -637,16 +643,14 @@ async def _async_device_config_resync(self) -> None:
 
     async def _async_connect(self) -> None:
         """Create connection."""
-        _, self._aio_protocol = await asyncio.wait_for(
-            self.loop.create_connection(  # type: ignore
+        async with asyncio_timeout(self.timeout):
+            _, self._aio_protocol = await self.loop.create_connection(  # type: ignore
                 lambda: AIOLEDENETProtocol(
                     self._async_data_recieved, self._async_connection_lost
                 ),
                 self.ipaddr,
                 self.port,
-            ),
-            timeout=self.timeout,
-        )
+            )
 
     def _async_connection_lost(self, exc: Optional[Exception]) -> None:
         """Called when the connection is lost."""
@@ -833,9 +837,8 @@ async def _async_determine_protocol(self) -> None:
                 self._determine_protocol_future = asyncio.Future()
                 self._aio_protocol.write(protocol.construct_state_query())
                 try:
-                    await asyncio.wait_for(
-                        self._determine_protocol_future, timeout=self.timeout
-                    )
+                    async with asyncio_timeout(self.timeout):
+                        await self._determine_protocol_future
                 except asyncio.TimeoutError:
                     if self._aio_protocol:
                         self._aio_protocol.close()

flux_led/aioscanner.py

@@ -4,6 +4,7 @@
 import time
 from typing import Callable, Dict, List, Optional, Tuple, cast
 
+from .aioutils import asyncio_timeout
 from .scanner import MESSAGE_SEND_INTERLEAVE_DELAY, BulbScanner, FluxLEDDiscovery
 
 _LOGGER = logging.getLogger(__name__)
@@ -68,7 +69,8 @@ async def _async_send_and_wait(
             events.append(event)
         for idx, command in enumerate(commands):
             self._send_message(transport, destination, command)
-            await asyncio.wait_for(event_map[idx].wait(), timeout=timeout)
+            async with asyncio_timeout(timeout):
+                await event_map[idx].wait()
 
     async def _async_send_commands_and_reboot(
         self,
@@ -122,9 +124,8 @@ async def _async_run_scan(
         time_out = timeout / self.BROADCAST_FREQUENCY
         while True:
             try:
-                await asyncio.wait_for(
-                    asyncio.shield(found_all_future), timeout=time_out
-                )
+                async with asyncio_timeout(time_out):
+                    await asyncio.shield(found_all_future)
             except asyncio.TimeoutError:
                 pass
             else:

flux_led/aioutils.py

@@ -0,0 +1,8 @@
+import sys
+
+__all__ = ["asyncio_timeout"]
+
+if sys.version_info[:2] < (3, 11):
+    from async_timeout import timeout as asyncio_timeout
+else:
+    from asyncio import timeout as asyncio_timeout

setup.py

@@ -31,7 +31,7 @@
     "wheel>=0.34.2",
 ]
 
-requirements = ["webcolors", 'typing_extensions;python_version<"3.8"']
+requirements = ["webcolors", 'typing_extensions;python_version<"3.8"', "async_timeout>=3.0.0"]
 
 
 extra_requirements = {