fix a windows specific issue where the wait_for_user_program_stop function doesn't detect the hub being turned off properly

Author: shaggysa

pybricksdev/cli/__init__.py

@@ -268,16 +268,18 @@ async def reconnect_hub():
             ]
             while True:
                 try:
-                    response = await hub.race_power_button_press(
-                        questionary.select(
-                            "Would you like to re-compile your code?",
-                            response_options,
-                            default=(
-                                response_options[0]
-                                if args.start
-                                else response_options[1]
-                            ),
-                        ).ask_async()
+                    response = await hub.race_disconnect(
+                        hub.race_power_button_press(
+                            questionary.select(
+                                "Would you like to re-compile your code?",
+                                response_options,
+                                default=(
+                                    response_options[0]
+                                    if args.start
+                                    else response_options[1]
+                                ),
+                            ).ask_async()
+                        )
                     )
                     with _get_script_path(args.file) as script_path:
                         if response == response_options[0]:
@@ -292,10 +294,22 @@ async def reconnect_hub():
                     # current program, so the menu was canceled and we are now printing
                     # the hub stdout until the user program ends on the hub.
                     try:
-                        await hub._wait_for_user_program_stop(5)
+                        await hub._wait_for_user_program_stop(
+                            2.25, raise_error_on_timeout=True
+                        )
+
                     except HubDisconnectError:
                         hub = await reconnect_hub()
 
+                    except asyncio.TimeoutError:
+                        # On windows, it takes significantly longer
+                        # for the device to register that the hub was
+                        # disconnected. If _wait_for_user_program_stop
+                        # throws a timeout error, we can assume that the
+                        # hub was disconnected.
+                        await hub.disconnect()
+                        hub = await reconnect_hub()
+
                 except HubDisconnectError:
                     # let terminal cool off before making a new prompt
                     await asyncio.sleep(0.3)

pybricksdev/connections/pybricks.py

@@ -717,24 +717,20 @@ async def race_power_button_press(self, awaitable: Awaitable[T]) -> T:
         power_button_press_event = asyncio.Event()
         power_button_press_task = asyncio.ensure_future(power_button_press_event.wait())
 
-        disconnect_event = asyncio.Event()
-        disconnect_task = asyncio.ensure_future(disconnect_event.wait())
-
-        def handle_disconnect(state: ConnectionState):
-            if state == ConnectionState.DISCONNECTED:
-                disconnect_event.set()
-
         def handle_power_button_press(status: StatusFlag):
             if status.value & StatusFlag.POWER_BUTTON_PRESSED:
                 power_button_press_event.set()
 
-        with self.status_observable.subscribe(
-            handle_power_button_press
-        ), self.connection_state_observable.subscribe(handle_disconnect):
-            done, pending = await asyncio.wait(
-                {awaitable_task, power_button_press_task, disconnect_task},
-                return_when=asyncio.FIRST_COMPLETED,
-            )
+        with self.status_observable.subscribe(handle_power_button_press):
+            try:
+                done, pending = await asyncio.wait(
+                    {awaitable_task, power_button_press_task},
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+            except BaseException:
+                awaitable_task.cancel()
+                power_button_press_task.cancel()
+                raise
 
         for t in pending:
             t.cancel()
@@ -743,11 +739,11 @@ def handle_power_button_press(status: StatusFlag):
             raise HubPowerButtonPressedError(
                 "the hub's power button was pressed during operation"
             )
-        if disconnect_task in done:
-            raise HubDisconnectError("the hub was disconnected during operation")
         return awaitable_task.result()
 
-    async def _wait_for_user_program_stop(self, program_start_timeout=1):
+    async def _wait_for_user_program_stop(
+        self, program_start_timeout=1, raise_error_on_timeout=False
+    ):
         user_program_running: asyncio.Queue[bool] = asyncio.Queue()
 
         with self.status_observable.pipe(
@@ -768,6 +764,8 @@ async def _wait_for_user_program_stop(self, program_start_timeout=1):
                         program_start_timeout,
                     )
                 except asyncio.TimeoutError:
+                    if raise_error_on_timeout:
+                        raise
                     # if it doesn't start, assume it was a very short lived
                     # program and we just missed the status message
                     logger.debug(