Merge pull request #1 from shaggysa/stay-connected

Fixes to the stay-connected feature

Author: shaggysa

pybricksdev/cli/__init__.py

@@ -20,6 +20,8 @@
 
 from pybricksdev import __name__ as MODULE_NAME
 from pybricksdev import __version__ as MODULE_VERSION
+from pybricksdev.ble.pybricks import StatusFlag
+from pybricksdev.connections import ConnectionState
 
 PROG_NAME = (
     f"{path.basename(sys.executable)} -m {MODULE_NAME}"
@@ -221,25 +223,95 @@ def is_pybricks_usb(dev):
         # Connect to the address and run the script
         await hub.connect()
         try:
-            while True:
-                with _get_script_path(args.file) as script_path:
-                    if args.start:
-                        await hub.run(script_path, args.wait)
-                    else:
-                        await hub.download(script_path)
-
-                if not args.wait or not args.stay_connected:
-                    break
-
-                resend = await questionary.select(
-                    "Would you like to resend your code?", choices=["Resend", "Exit"]
-                ).ask_async()
-
-                if resend == "Exit":
-                    break
-
-        except RuntimeError:
-            print("The hub is no longer connected.")
+            with _get_script_path(args.file) as script_path:
+                if args.start:
+                    await hub.run(script_path, args.wait or args.stay_connected)
+                else:
+                    if args.stay_connected:
+                        hub.print_output = True
+                        hub._enable_line_handler = True
+                    await hub.download(script_path)
+
+            if args.stay_connected:
+                response_options = [
+                    "Recompile and Run",
+                    "Recompile and Download",
+                    "Exit",
+                ]
+                while True:
+                    try:
+                        response = await hub.race_power_button_press(
+                            questionary.select(
+                                "Would you like to re-compile your code?",
+                                response_options,
+                            ).ask_async()
+                        )
+                    except RuntimeError as e:
+
+                        async def reconnect_hub():
+                            if await questionary.confirm(
+                                "\nThe hub has been disconnected. Would you like to re-connect?"
+                            ).ask_async():
+                                if args.conntype == "ble":
+                                    print(
+                                        f"Searching for {args.name or 'any hub with Pybricks service'}..."
+                                    )
+                                    device_or_address = await find_ble(args.name)
+                                    hub = PybricksHubBLE(device_or_address)
+                                elif args.conntype == "usb":
+                                    device_or_address = find_usb(
+                                        custom_match=is_pybricks_usb
+                                    )
+                                    hub = PybricksHubUSB(device_or_address)
+
+                                await hub.connect()
+                                # re-enable echoing of the hub's stdout
+                                hub.print_output = True
+                                hub._enable_line_handler = True
+                                return hub
+
+                            else:
+                                exit()
+
+                        if (
+                            hub.status_observable.value
+                            & StatusFlag.POWER_BUTTON_PRESSED
+                        ):
+                            try:
+                                await hub._wait_for_user_program_stop(2.1)
+                                continue
+
+                            except RuntimeError as e:
+                                if (
+                                    hub.connection_state_observable.value
+                                    == ConnectionState.DISCONNECTED
+                                ):
+                                    hub = await reconnect_hub()
+                                    continue
+
+                                else:
+                                    raise e
+
+                        elif (
+                            hub.connection_state_observable.value
+                            == ConnectionState.DISCONNECTED
+                        ):
+                            # let terminal cool off before making a new prompt
+                            await asyncio.sleep(0.3)
+
+                            hub = await reconnect_hub()
+                            continue
+
+                        else:
+                            raise e
+
+                    with _get_script_path(args.file) as script_path:
+                        if response == response_options[0]:
+                            await hub.run(script_path, True)
+                        elif response == response_options[1]:
+                            await hub.download(script_path)
+                        else:
+                            exit(1)
 
         finally:
             await hub.disconnect()

pybricksdev/connections/pybricks.py

@@ -678,7 +678,60 @@ async def send_block(data: bytes) -> None:
         if wait:
             await self._wait_for_user_program_stop()
 
-    async def _wait_for_user_program_stop(self):
+    async def race_power_button_press(self, awaitable: Awaitable[T]) -> T:
+        """
+        Races an awaitable against the user pressing the power button of the hub.
+
+        If the power button is pressed or the hub becomes disconnected before the awaitable is complete, a
+        ``RuntimeError`` is raised and the awaitable is canceled.
+
+        Otherwise, the result of the awaitable is returned. If the awaitable
+        raises an exception, that exception will be raised.
+
+        Args:
+            awaitable: Any awaitable such as a coroutine.
+
+        Returns:
+            The result of the awaitable.
+
+        Raises:
+            RuntimeError:
+                Thrown if the hub's power button is pressed or the hub is disconnected.
+        """
+        awaitable_task = asyncio.ensure_future(awaitable)
+
+        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
+        ) and 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,
+            )
+
+        for t in pending:
+            t.cancel()
+
+        if power_button_press_task in done:
+            raise RuntimeError("the hub's power button was pressed during operation")
+        elif disconnect_task in done:
+            raise RuntimeError("the hub was disconnected during operation")
+        return awaitable_task.result()
+
+    async def _wait_for_user_program_stop(self, program_start_timeout=1):
         user_program_running: asyncio.Queue[bool] = asyncio.Queue()
 
         with self.status_observable.pipe(
@@ -695,7 +748,8 @@ async def _wait_for_user_program_stop(self):
                 # for it to start
                 try:
                     await asyncio.wait_for(
-                        self.race_disconnect(user_program_running.get()), 1
+                        self.race_disconnect(user_program_running.get()),
+                        program_start_timeout,
                     )
                 except asyncio.TimeoutError:
                     # if it doesn't start, assume it was a very short lived