Add the abilty to resend code though bluetooth in the same session

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

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"

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

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"

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(