minor fixes and cleanup

Author: shaggysa

pybricksdev/cli/__init__.py

@@ -230,75 +230,77 @@ def is_pybricks_usb(dev):
                     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 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()
+            if not args.stay_connected:
+                exit()
+
+            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:
+                    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, wait=True)
+                        elif response == response_options[1]:
+                            await hub.download(script_path)
+                        else:
+                            exit()
+
+                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_user_program_stop(5)
                     except HubDisconnectError:
-                        # let terminal cool off before making a new prompt
-                        await asyncio.sleep(0.3)
                         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()
 

pybricksdev/connections/pybricks.py

@@ -65,12 +65,12 @@
 T = TypeVar("T")
 
 
-class HubDisconnectError(Exception):
+class HubDisconnectError(RuntimeError):
     """Raise when a hub disconnect occurs."""
 
 
-class HubPowerButtonPressedError(Exception):
-    """Raise when the hub's power button is pressed."""
+class HubPowerButtonPressedError(RuntimeError):
+    """Raise when a task was canceled because the hub's power button was pressed."""
 
 
 class PybricksHub:
@@ -689,12 +689,13 @@ async def send_block(data: bytes) -> None:
     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.
-
+        :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.
@@ -703,8 +704,8 @@ async def race_power_button_press(self, awaitable: Awaitable[T]) -> T:
             The result of the awaitable.
 
         Raises:
-            RuntimeError:
-                Thrown if the hub's power button is pressed or the hub is disconnected.
+            HubPowerButtonPressedError
+            HubDisconnectError
         """
         awaitable_task = asyncio.ensure_future(awaitable)
 
@@ -724,7 +725,7 @@ def handle_power_button_press(status: StatusFlag):
 
         with self.status_observable.subscribe(
             handle_power_button_press
-        ) and self.connection_state_observable.subscribe(handle_disconnect):
+        ), 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,
@@ -737,7 +738,7 @@ def handle_power_button_press(status: StatusFlag):
             raise HubPowerButtonPressedError(
                 "the hub's power button was pressed during operation"
             )
-        elif disconnect_task in done:
+        if disconnect_task in done:
             raise HubDisconnectError("the hub was disconnected during operation")
         return awaitable_task.result()