cli: add --stay-connected argument run command

shaggysa · web-flow · commit 585168c7015e · 2025-09-10T21:22:00.000-05:00
Add a `--stay-connected` argument that causes the `pybricksdev run` command to stay connected after the user program ends (implies `--wait`). After this, an interactive menu is show to allow compiling and running or downloading the program again. Also cancels the menu if the program is started by pressing the button on the hub so that stdout can print.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Added the `--stay-connected` arg to the `pybricksdev run`
+    command, allowing re-compiling and running the input file. Also echoes
+    the hub's output to the console when manually running a program.
+    ([pybricksdev#122])
+
+[pybricksdev#122]: https://github.com/pybricks/pybricksdev/pull/122
+
+
 ## [2.0.1] - 2025-08-11
 
 ### Fixed
diff --git a/poetry.lock b/poetry.lock
diff --git a/pybricksdev/cli/__init__.py b/pybricksdev/cli/__init__.py
@@ -15,10 +15,15 @@
 from typing import ContextManager, TextIO
 
 import argcomplete
+import questionary
 from argcomplete.completers import FilesCompleter
 
 from pybricksdev import __name__ as MODULE_NAME
 from pybricksdev import __version__ as MODULE_VERSION
+from pybricksdev.connections.pybricks import (
+    HubDisconnectError,
+    HubPowerButtonPressedError,
+)
 
 PROG_NAME = (
     f"{path.basename(sys.executable)} -m {MODULE_NAME}"
@@ -160,6 +165,13 @@ def add_parser(self, subparsers: argparse._SubParsersAction):
             default=True,
         )
 
+        parser.add_argument(
+            "--stay-connected",
+            help="Add a menu option to resend the code with bluetooth instead of disconnecting from the robot after the program ends.",
+            action=argparse.BooleanOptionalAction,
+            default=False,
+        )
+
     async def run(self, args: argparse.Namespace):
 
         # Pick the right connection
@@ -215,14 +227,99 @@ def is_pybricks_usb(dev):
         try:
             with _get_script_path(args.file) as script_path:
                 if args.start:
-                    await hub.run(script_path, args.wait)
+                    await hub.run(script_path, args.wait or args.stay_connected)
                 else:
+                    if args.stay_connected:
+                        # if the user later starts the program by pressing the button on the hub,
+                        # we still want the hub stdout to print to Python's stdout
+                        hub.print_output = True
+                        hub._enable_line_handler = True
                     await hub.download(script_path)
+
+            if not args.stay_connected:
+                return
+
+            async def reconnect_hub():
+                if not await questionary.confirm(
+                    "\nThe hub has been disconnected. Would you like to re-connect?"
+                ).ask_async():
+                    exit()
+
+                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._enable_line_handler = True
+                hub.print_output = True
+                return hub
+
+            response_options = [
+                "Recompile and Run",
+                "Recompile and Download",
+                "Exit",
+            ]
+            while True:
+                try:
+                    if args.file is sys.stdin:
+                        await hub.race_disconnect(
+                            hub.race_power_button_press(
+                                questionary.press_any_key_to_continue(
+                                    "The hub will stay connected and echo its output to the terminal. Press any key to exit."
+                                ).ask_async()
+                            )
+                        )
+                        return
+                    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]:
+                            await hub.run(script_path, wait=True)
+                        elif response == response_options[1]:
+                            await hub.download(script_path)
+                        else:
+                            return
+
+                except HubPowerButtonPressedError:
+                    # This means the user pressed the button on the hub to re-start the
+                    # 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_power_button_release()
+                        await hub._wait_for_user_program_stop()
+
+                    except HubDisconnectError:
+                        hub = await reconnect_hub()
+
+                except HubDisconnectError:
+                    # let terminal cool off before making a new prompt
+                    await asyncio.sleep(0.3)
+                    hub = await reconnect_hub()
+
         finally:
             await hub.disconnect()
 
 
 class Flash(Tool):
+
     def add_parser(self, subparsers: argparse._SubParsersAction):
         parser = subparsers.add_parser(
             "flash", help="flash firmware on a LEGO Powered Up device"
diff --git a/pybricksdev/connections/pybricks.py b/pybricksdev/connections/pybricks.py
@@ -65,6 +65,14 @@
 T = TypeVar("T")
 
 
+class HubDisconnectError(RuntimeError):
+    """Raise when a hub disconnect occurs."""
+
+
+class HubPowerButtonPressedError(RuntimeError):
+    """Raise when a task was canceled because the hub's power button was pressed."""
+
+
 class PybricksHub:
     EOL = b"\r\n"  # MicroPython EOL
 
@@ -363,7 +371,7 @@ def handle_disconnect(state: ConnectionState):
                 t.cancel()
 
             if awaitable_task not in done:
-                raise RuntimeError("disconnected during operation")
+                raise HubDisconnectError("disconnected during operation")
 
             return awaitable_task.result()
 
@@ -683,6 +691,92 @@ async def send_block(data: bytes) -> None:
         if wait:
             await self._wait_for_user_program_stop()
 
+    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
+        :class:`HubPowerButtonPressedError` or :class:`HubDisconnectError` 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.
+        The intended purpose of this function is to detect when
+        the user manually starts a program on the hub. It is used instead of the program running flag
+        because the hub can send info through stdout before we can detect a change in the program running flag.
+
+        Args:
+            awaitable: Any awaitable such as a coroutine.
+
+        Returns:
+            The result of the awaitable.
+
+        Raises:
+            HubPowerButtonPressedError
+            HubDisconnectError
+        """
+        awaitable_task = asyncio.ensure_future(awaitable)
+
+        power_button_press_event = asyncio.Event()
+        power_button_press_task = asyncio.ensure_future(power_button_press_event.wait())
+
+        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):
+            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()
+
+        if power_button_press_task in done:
+            raise HubPowerButtonPressedError(
+                "the hub's power button was pressed during operation"
+            )
+        return awaitable_task.result()
+
+    async def _wait_for_power_button_release(self):
+        power_button_pressed: asyncio.Queue[bool] = asyncio.Queue()
+
+        with self.status_observable.pipe(
+            op.map(lambda s: bool(s & StatusFlag.POWER_BUTTON_PRESSED)),
+            op.distinct_until_changed(),
+        ).subscribe(lambda s: power_button_pressed.put_nowait(s)):
+            # The first item in the queue is the current status. The status
+            # could change before or after the last checksum is received,
+            # so this could be either true or false.
+            is_pressed = await self.race_disconnect(power_button_pressed.get())
+
+            if not is_pressed:
+                # If the button isn't already pressed,
+                # wait a short time for it to become pressed
+                try:
+                    await asyncio.wait_for(
+                        self.race_disconnect(power_button_pressed.get()),
+                        1,
+                    )
+                except asyncio.TimeoutError:
+                    # If the button never shows as pressed,
+                    # assume that we just missed the status flag
+                    logger.debug(
+                        "timed out waiting for power button press, assuming is was a short press"
+                    )
+                    return
+
+            # At this point, we know the button is pressed, so the
+            # next item in the queue must indicate the button has
+            # been released.
+            is_pressed = await self.race_disconnect(power_button_pressed.get())
+
+            # maybe catch mistake if the code is changed
+            assert not is_pressed
+
     async def _wait_for_user_program_stop(self):
         user_program_running: asyncio.Queue[bool] = asyncio.Queue()
 
@@ -700,7 +794,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()),
+                        1,
                     )
                 except asyncio.TimeoutError:
                     # if it doesn't start, assume it was a very short lived
diff --git a/pyproject.toml b/pyproject.toml
@@ -48,6 +48,7 @@ typing-extensions = ">=4.3.0"
 reactivex = {version = ">=4.0.4", python = "<4"}
 hidapi = ">=0.14.0"
 pybricks = {version = ">=3", allow-prereleases = true, python = "<4"}
+questionary = {version = ">=2.1.1", python = "<4"}
 
 [tool.poetry.group.lint.dependencies]
 black = ">=23,<25"
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -86,6 +86,7 @@ async def test_download_ble(self):
                 name="MyHub",
                 start=False,
                 wait=False,
+                stay_connected=False,
             )
 
             mock_hub_class = stack.enter_context(
@@ -135,6 +136,7 @@ async def test_download_usb(self):
                 name=None,
                 start=False,
                 wait=False,
+                stay_connected=False,
             )
 
             mock_hub_class = stack.enter_context(
@@ -175,6 +177,7 @@ async def test_download_stdin(self):
             name="MyHub",
             start=False,
             wait=False,
+            stay_connected=False,
         )
 
         # Set up mocks using ExitStack
@@ -227,6 +230,7 @@ async def test_download_connection_error(self):
                 name="MyHub",
                 start=False,
                 wait=False,
+                stay_connected=False,
             )
 
             stack.enter_context(
@@ -273,6 +277,7 @@ async def test_run_ble(self):
                 name="MyHub",
                 start=True,
                 wait=True,
+                stay_connected=False,
             )
 
             mock_hub_class = stack.enter_context(
@@ -321,6 +326,7 @@ async def test_run_usb(self):
                 name=None,
                 start=True,
                 wait=True,
+                stay_connected=False,
             )
 
             mock_hub_class = stack.enter_context(
@@ -360,6 +366,7 @@ async def test_run_stdin(self):
             name="MyHub",
             start=True,
             wait=True,
+            stay_connected=False,
         )
 
         # Set up mocks using ExitStack
@@ -414,6 +421,7 @@ async def test_run_connection_error(self):
                 name="MyHub",
                 start=False,
                 wait=True,
+                stay_connected=False,
             )
 
             stack.enter_context(
