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

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,84 @@ 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:
+                        hub.print_output = True
+                        hub._enable_line_handler = True
                     await hub.download(script_path)
+
+            if args.stay_connected:
+
+                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.log_file = None
+                    hub.output = []
+                    hub._stdout_buf.clear()
+                    hub._stdout_line_queue = asyncio.Queue()
+                    hub.print_output = True
+                    hub._enable_line_handler = True
+                    return hub
+
+                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,
+                                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, True)
+                            elif response == response_options[1]:
+                                await hub.download(script_path)
+                            else:
+                                exit()
+
+                    except HubPowerButtonPressedError:
+                        try:
+                            await hub._wait_for_user_program_stop(5)
+                        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(Exception):
+    """Raise when a hub disconnect occurs."""
+
+
+class HubPowerButtonPressedError(Exception):
+    """Raise when the hub's power button is pressed."""
+
+
 class PybricksHub:
     EOL = b"\r\n"  # MicroPython EOL
 
@@ -358,7 +366,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()
 
@@ -678,7 +686,62 @@ 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 HubPowerButtonPressedError(
+                "the hub's power button was pressed during operation"
+            )
+        elif 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):
         user_program_running: asyncio.Queue[bool] = asyncio.Queue()
 
         with self.status_observable.pipe(
@@ -695,7 +758,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

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(