diff --git a/CHANGELOG.md b/CHANGELOG.md index 48bcb4d..bb159f7 100644 --- 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 index c1c03c0..017c5bc 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1079,6 +1079,22 @@ files = [ {file = "pyusb-1.3.1.tar.gz", hash = "sha256:3af070b607467c1c164f49d5b0caabe8ac78dbed9298d703a8dbf9df4052d17e"}, ] +[[package]] +name = "questionary" +version = "2.1.1" +description = "Python library to build pretty command line user prompts ⭐️" +optional = false +python-versions = ">=3.9" +groups = ["main"] +markers = "python_version < \"4\"" +files = [ + {file = "questionary-2.1.1-py3-none-any.whl", hash = "sha256:a51af13f345f1cdea62347589fbb6df3b290306ab8930713bfae4d475a7d4a59"}, + {file = "questionary-2.1.1.tar.gz", hash = "sha256:3d7e980292bb0107abaa79c68dd3eee3c561b83a0f89ae482860b181c8bd412d"}, +] + +[package.dependencies] +prompt_toolkit = ">=2.0,<4.0" + [[package]] name = "reactivex" version = "4.0.4" @@ -1716,4 +1732,4 @@ all = ["winrt-Windows.Foundation.Collections[all] (>=3.2.1.0,<3.3.0.0)", "winrt- [metadata] lock-version = "2.1" python-versions = ">=3.10" -content-hash = "3d2a87a44fe61f7cbe48f13302f75cdcfa991b2ffc546e12818b9999050c5959" +content-hash = "71dd6b61b42043573f4f6f7d1b89ae384c60f320c10d45b9095125d1dec015a1" diff --git a/pybricksdev/cli/__init__.py b/pybricksdev/cli/__init__.py index a2383c0..a499bd3 100644 --- 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 index 1eb2661..709f238 100644 --- 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 index ae1510c..37aa9c7 100644 --- 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 index 28f2330..b222405 100644 --- 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(