20311043488

Committed 17 Dec 2025 05:08PM UTC coverage: 55.368% (+1.5%) from 53.832%

Build # 20311043488

Build Type

Pull #126

github

Committed by

web-flow

Commit Message

Merge 884adce5f into a1a9a71de

Pull Request Pull Request #126: Catch a syntax error when parsing an input file with the --stay-connected flag active

Run Details

82 of 264 branches covered (31.06%)

Branch coverage included in aggregate %.

2110 of 3695 relevant lines covered (57.1%)

0.57 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

60.83

pybricksdev/cli/__init__.py

# SPDX-License-Identifier: MIT
# Copyright (c) 2019-2024 The Pybricks Authors

"""Command line wrapper around pybricksdev library."""

import argparse
import asyncio
import contextlib
import logging
import os
import subprocess
import sys
from abc import ABC, abstractmethod
from enum import IntEnum
from os import PathLike, path
from tempfile import NamedTemporaryFile
from typing import ContextManager, TextIO

import argcomplete
import questionary
from argcomplete.completers import FilesCompleter
from packaging.version import Version

from pybricksdev import __name__ as MODULE_NAME
from pybricksdev import __version__ as MODULE_VERSION
from pybricksdev.connections.pybricks import (
    HubDisconnectError,
    HubPowerButtonPressedError,
    PybricksHub,
)

PROG_NAME = (
    f"{path.basename(sys.executable)} -m {MODULE_NAME}"
    if sys.argv[0].endswith("__main__.py")
    else path.basename(sys.argv[0])
)


class Tool(ABC):
    """Common base class for tool implementations."""

    @abstractmethod
    def add_parser(self, subparsers: argparse._SubParsersAction):
        """
        Overriding methods must at least do the following::

            parser = subparsers.add_parser('tool', ...)
            parser.tool = self

        Then additional arguments can be added using the ``parser`` object.
        """
        pass

    @abstractmethod
    async def run(self, args: argparse.Namespace):
        """
        Overriding methods should provide an implementation to run the tool.
        """
        pass


def _get_script_path(file: TextIO) -> ContextManager[PathLike]:
    """
    Gets the path to a script on the file system.

    If the file is ``sys.stdin``, the contents are copied to a temporary file
    and the path to the temporary file is returned. Otherwise, the file is closed
    and the path is returned.

    The context manager will delete the temporary file, if applicable.
    """
    if file is sys.stdin:
        # Have to close the temp file so that mpy-cross can read it, so we
        # create our own context manager to delete the file when we are done
        # using it.

        @contextlib.contextmanager
        def temp_context():
            try:
                with NamedTemporaryFile("wb", suffix=".py", delete=False) as temp:
                    temp.write(file.buffer.read())

                yield temp.name
            finally:
                try:
                    os.remove(temp.name)
                except NameError:
                    # if NamedTemporaryFile() throws, temp is not defined
                    pass
                except OSError:
                    # file was already deleted or other strangeness
                    pass

        return temp_context()

    file.close()
    return contextlib.nullcontext(file.name)


class Compile(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        parser = subparsers.add_parser(
            "compile",
            help="compile a Pybricks program without running it",
        )
        parser.add_argument(
            "file",
            metavar="<file>",
            help="path to a MicroPython script or `-` for stdin",
            type=argparse.FileType(encoding="utf-8"),
        )
        parser.add_argument(
            "--abi",
            metavar="<n>",
            help="the MPY ABI version, one of %(choices)s (default: %(default)s)",
            default=6,
            choices=[5, 6],
            type=int,
        )
        parser.add_argument(
            "--bin",
            action="store_true",
            help="output unformatted binary data only (useful for pipes)",
        )
        parser.tool = self

    async def run(self, args: argparse.Namespace):
        from pybricksdev.compile import compile_multi_file, print_mpy

        with _get_script_path(args.file) as script_path:
            mpy = await compile_multi_file(script_path, args.abi)
        if args.bin:
            sys.stdout.buffer.write(mpy)
            return
        print_mpy(mpy)


class Run(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        parser = subparsers.add_parser(
            "run",
            help="run a Pybricks program",
        )
        parser.tool = self
        parser.add_argument(
            "conntype",
            metavar="<connection type>",
            help="connection type: %(choices)s",
            choices=["ble", "usb"],
        )
        parser.add_argument(
            "file",
            metavar="<file>",
            help="path to a MicroPython script or `-` for stdin",
            type=argparse.FileType(encoding="utf-8"),
        )
        parser.add_argument(
            "-n",
            "--name",
            metavar="<name>",
            required=False,
            help="Bluetooth device name or Bluetooth address for BLE connection; "
            "serial port name for USB connection",
        )

        parser.add_argument(
            "--start",
            help="Start the program immediately after downloading it.",
            action=argparse.BooleanOptionalAction,
            default=True,
        )

        parser.add_argument(
            "--wait",
            help="Wait for the program to complete before disconnecting. Only applies when starting program right away.",
            action=argparse.BooleanOptionalAction,
            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 stay_connected_menu(self, hub: PybricksHub, args: argparse.Namespace):

        if args.conntype == "ble":
            from pybricksdev.ble import find_device as find_ble
            from pybricksdev.connections.pybricks import PybricksHubBLE
        else:
            from usb.core import find as find_usb

            from pybricksdev.connections.pybricks import PybricksHubUSB
            from pybricksdev.usb import (
                EV3_USB_PID,
                LEGO_USB_VID,
                MINDSTORMS_INVENTOR_USB_PID,
                NXT_USB_PID,
                SPIKE_ESSENTIAL_USB_PID,
                SPIKE_PRIME_USB_PID,
            )

            def is_pybricks_usb(dev):
                return (
                    (dev.idVendor == LEGO_USB_VID)
                    and (
                        dev.idProduct
                        in [
                            NXT_USB_PID,
                            EV3_USB_PID,
                            SPIKE_PRIME_USB_PID,
                            SPIKE_ESSENTIAL_USB_PID,
                            MINDSTORMS_INVENTOR_USB_PID,
                        ]
                    )
                    and dev.product.endswith("Pybricks")
                )

        class ResponseOptions(IntEnum):
            RECOMPILE_RUN = 0
            RECOMPILE_DOWNLOAD = 1
            RUN_STORED = 2
            CHANGE_TARGET_FILE = 3
            EXIT = 4

        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",
            "Run Stored Program",
            "Change Target File",
            "Exit",
        ]
        # the entry that is selected by default when the menu opens
        # this is overridden after the user picks an option
        # so that the default option is always the one that was last chosen
        default_response_option = (
            ResponseOptions.RECOMPILE_RUN
            if args.start
            else ResponseOptions.RECOMPILE_DOWNLOAD
        )

        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(
                            f"Would you like to re-compile {os.path.basename(args.file.name)}?",
                            response_options,
                            default=(response_options[default_response_option]),
                        ).ask_async()
                    )
                )

                default_response_option = response_options.index(response)

                match response_options.index(response):

                    case ResponseOptions.RECOMPILE_RUN:
                        with _get_script_path(args.file) as script_path:
                            await hub.run(script_path, wait=True)

                    case ResponseOptions.RECOMPILE_DOWNLOAD:
                        with _get_script_path(args.file) as script_path:
                            await hub.download(script_path)

                    case ResponseOptions.RUN_STORED:
                        if hub.fw_version < Version("3.2.0-beta.4"):
                            print(
                                "Running a stored program remotely is only supported in the hub firmware version >= v3.2.0."
                            )
                        else:
                            await hub.start_user_program()
                            await hub._wait_for_user_program_stop()

                    case ResponseOptions.CHANGE_TARGET_FILE:
                        args.file.close()
                        while True:
                            try:
                                args.file = open(
                                    await hub.race_disconnect(
                                        hub.race_power_button_press(
                                            questionary.path(
                                                "What file would you like to use?"
                                            ).ask_async()
                                        )
                                    ),
                                    encoding="utf-8",
                                )
                                break
                            except FileNotFoundError:
                                print("The file was not found. Please try again.")
                        # send the new target file to the hub
                        with _get_script_path(args.file) as script_path:
                            await hub.download(script_path)

                    case _:
                        return

            except subprocess.CalledProcessError as e:
                print()
                print("mpy-cross failed to compile the program:")
                print(e.stderr.decode())

            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()

    async def run(self, args: argparse.Namespace):

        # Pick the right connection
        if args.conntype == "ble":
            from pybricksdev.ble import find_device as find_ble
            from pybricksdev.connections.pybricks import PybricksHubBLE

            # It is a Pybricks Hub with BLE. Device name or address is given.
            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":
            from usb.core import find as find_usb

            from pybricksdev.connections.pybricks import PybricksHubUSB
            from pybricksdev.usb import (
                EV3_USB_PID,
                LEGO_USB_VID,
                MINDSTORMS_INVENTOR_USB_PID,
                NXT_USB_PID,
                SPIKE_ESSENTIAL_USB_PID,
                SPIKE_PRIME_USB_PID,
            )

            def is_pybricks_usb(dev):
                return (
                    (dev.idVendor == LEGO_USB_VID)
                    and (
                        dev.idProduct
                        in [
                            NXT_USB_PID,
                            EV3_USB_PID,
                            SPIKE_PRIME_USB_PID,
                            SPIKE_ESSENTIAL_USB_PID,
                            MINDSTORMS_INVENTOR_USB_PID,
                        ]
                    )
                    and dev.product.endswith("Pybricks")
                )

            device_or_address = find_usb(custom_match=is_pybricks_usb)

            if device_or_address is None:
                print("Pybricks Hub not found.", file=sys.stderr)
                exit(1)

            hub = PybricksHubUSB(device_or_address)
        else:
            raise ValueError(f"Unknown connection type: {args.conntype}")

        # Connect to the address and run the script
        await hub.connect()
        try:
            with _get_script_path(args.file) as script_path:
                if args.start:
                    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)

        except subprocess.CalledProcessError as e:
            print()
            print("mpy-cross failed to compile the program:")
            print(e.stderr.decode())

        finally:
            if args.stay_connected:
                await self.stay_connected_menu(hub, args)
            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"
        )
        parser.tool = self

        parser.add_argument(
            "firmware",
            metavar="<firmware-file>",
            type=argparse.FileType(mode="rb"),
            help="the firmware .zip file",
        ).completer = FilesCompleter(allowednames=(".zip",))

        parser.add_argument(
            "-n", "--name", metavar="<name>", type=str, help="a custom name for the hub"
        )

    def run(self, args: argparse.Namespace):
        from pybricksdev.cli.flash import flash_firmware

        return flash_firmware(args.firmware, args.name)


class DFUBackup(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        parser = subparsers.add_parser("backup", help="backup firmware using DFU")
        parser.tool = self
        parser.add_argument(
            "firmware",
            metavar="<firmware-file>",
            type=argparse.FileType(mode="wb"),
            help="the firmware .bin file",
        ).completer = FilesCompleter(allowednames=(".bin",))

    async def run(self, args: argparse.Namespace):
        from pybricksdev.dfu import backup_dfu

        backup_dfu(args.firmware)


class DFURestore(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        parser = subparsers.add_parser(
            "restore",
            help="restore firmware using DFU",
        )
        parser.tool = self
        parser.add_argument(
            "firmware",
            metavar="<firmware-file>",
            type=argparse.FileType(mode="rb"),
            help="the firmware .bin file",
        ).completer = FilesCompleter(allowednames=(".bin",))

    async def run(self, args: argparse.Namespace):
        from pybricksdev.dfu import restore_dfu

        restore_dfu(args.firmware)


class DFU(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        self.parser = subparsers.add_parser(
            "dfu",
            help="use DFU to backup or restore firmware",
        )
        self.parser.tool = self
        self.subparsers = self.parser.add_subparsers(
            metavar="<action>", dest="action", help="the action to perform"
        )

        for tool in DFUBackup(), DFURestore():
            tool.add_parser(self.subparsers)

    def run(self, args: argparse.Namespace):
        if args.action not in self.subparsers.choices:
            self.parser.error(
                f'Missing name of action: {"|".join(self.subparsers.choices.keys())}'
            )

        return self.subparsers.choices[args.action].tool.run(args)


class OADFlash(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        parser = subparsers.add_parser(
            "flash",
            help="update firmware on a LEGO Powered Up device using TI OAD",
        )
        parser.tool = self
        parser.add_argument(
            "firmware",
            metavar="<firmware-file>",
            type=argparse.FileType(mode="rb"),
            help="the firmware .oda file",
        ).completer = FilesCompleter(allowednames=(".oda",))

    async def run(self, args: argparse.Namespace):
        from pybricksdev.cli.oad import flash_oad_image

        await flash_oad_image(args.firmware)


class OADInfo(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        parser = subparsers.add_parser(
            "info",
            help="get information about firmware on a LEGO Powered Up device using TI OAD",
        )
        parser.tool = self

    async def run(self, args: argparse.Namespace):
        from pybricksdev.cli.oad import dump_oad_info

        await dump_oad_info()


class OAD(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        self.parser = subparsers.add_parser(
            "oad",
            help="update firmware on a LEGO Powered Up device using TI OAD",
        )
        self.parser.tool = self
        self.subparsers = self.parser.add_subparsers(
            metavar="<action>", dest="action", help="the action to perform"
        )

        for tool in OADFlash(), OADInfo():
            tool.add_parser(self.subparsers)

    def run(self, args: argparse.Namespace):
        if args.action not in self.subparsers.choices:
            self.parser.error(
                f'Missing name of action: {"|".join(self.subparsers.choices.keys())}'
            )

        return self.subparsers.choices[args.action].tool.run(args)


class LWP3Repl(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        parser = subparsers.add_parser(
            "repl",
            help="interactive REPL for sending and receiving LWP3 messages",
        )
        parser.tool = self

    def run(self, args: argparse.Namespace):
        from pybricksdev.cli.lwp3.repl import repl, setup_repl_logging

        setup_repl_logging()
        return repl()


class LWP3(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        self.parser = subparsers.add_parser(
            "lwp3", help="interact with devices using LWP3"
        )
        self.parser.tool = self
        self.subparsers = self.parser.add_subparsers(
            metavar="<lwp3-tool>", dest="lwp3_tool", help="the tool to run"
        )

        for tool in (LWP3Repl(),):
            tool.add_parser(self.subparsers)

    def run(self, args: argparse.Namespace):
        if args.lwp3_tool not in self.subparsers.choices:
            self.parser.error(
                f'Missing name of tool: {"|".join(self.subparsers.choices.keys())}'
            )

        return self.subparsers.choices[args.lwp3_tool].tool.run(args)


class Udev(Tool):
    def add_parser(self, subparsers: argparse._SubParsersAction):
        parser = subparsers.add_parser("udev", help="print udev rules to stdout")
        parser.tool = self

    async def run(self, args: argparse.Namespace):
        from importlib.resources import read_text

        from pybricksdev import resources

        print(read_text(resources, resources.UDEV_RULES))


def main():
    """Runs ``pybricksdev`` command line interface."""

    if sys.platform == "win32":
        # Hack around bad side-effects of pythoncom on Windows
        try:
            from bleak_winrt._winrt import MTA, init_apartment
        except ImportError:
            from winrt._winrt import MTA, init_apartment

        init_apartment(MTA)

    # Provide main description and help.
    parser = argparse.ArgumentParser(
        prog=PROG_NAME,
        description="Utilities for Pybricks developers.",
        epilog="Run `%(prog)s <tool> --help` for tool-specific arguments.",
    )

    parser.add_argument(
        "-v", "--version", action="version", version=f"{MODULE_NAME} v{MODULE_VERSION}"
    )
    parser.add_argument(
        "-d", "--debug", action="store_true", help="enable debug logging"
    )

    subparsers = parser.add_subparsers(
        metavar="<tool>",
        dest="tool",
        help="the tool to use",
    )

    for tool in Compile(), Run(), Flash(), DFU(), OAD(), LWP3(), Udev():
        tool.add_parser(subparsers)

    argcomplete.autocomplete(parser)
    args = parser.parse_args()

    logging.basicConfig(
        format="%(asctime)s: %(levelname)s: %(name)s: %(message)s",
        level=logging.DEBUG if args.debug else logging.WARNING,
    )

    if not args.tool:
        parser.error(f'Missing name of tool: {"|".join(subparsers.choices.keys())}')

    asyncio.run(subparsers.choices[args.tool].tool.run(args))

1	# SPDX-License-Identifier: MIT
2	# Copyright (c) 2019-2024 The Pybricks Authors
3
4	"""Command line wrapper around pybricksdev library."""
5
6	import argparse	1✔
7	import asyncio	1✔
8	import contextlib	1✔
9	import logging	1✔
10	import os	1✔
11	import subprocess	1✔
12	import sys	1✔
13	from abc import ABC, abstractmethod	1✔
14	from enum import IntEnum	1✔
15	from os import PathLike, path	1✔
16	from tempfile import NamedTemporaryFile	1✔
17	from typing import ContextManager, TextIO	1✔
18
19	import argcomplete	1✔
20	import questionary	1✔
21	from argcomplete.completers import FilesCompleter	1✔
22	from packaging.version import Version	1✔
23
24	from pybricksdev import __name__ as MODULE_NAME	1✔
25	from pybricksdev import __version__ as MODULE_VERSION	1✔
26	from pybricksdev.connections.pybricks import (	1✔
27	HubDisconnectError,
28	HubPowerButtonPressedError,
29	PybricksHub,
30	)
31
32	PROG_NAME = (	1✔
33	f"{path.basename(sys.executable)} -m {MODULE_NAME}"
34	if sys.argv[0].endswith("__main__.py")
35	else path.basename(sys.argv[0])
36	)
37
38
39	class Tool(ABC):	1✔
40	"""Common base class for tool implementations."""
41
42	@abstractmethod	1✔
43	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
44	"""
45	Overriding methods must at least do the following::
46
47	parser = subparsers.add_parser('tool', ...)
48	parser.tool = self
49
50	Then additional arguments can be added using the ``parser`` object.
51	"""
52	pass	×
53
54	@abstractmethod	1✔
55	async def run(self, args: argparse.Namespace):	1✔
56	"""
57	Overriding methods should provide an implementation to run the tool.
58	"""
59	pass	×
60
61
62	def _get_script_path(file: TextIO) -> ContextManager[PathLike]:	1✔
63	"""
64	Gets the path to a script on the file system.
65
66	If the file is ``sys.stdin``, the contents are copied to a temporary file
67	and the path to the temporary file is returned. Otherwise, the file is closed
68	and the path is returned.
69
70	The context manager will delete the temporary file, if applicable.
71	"""
72	if file is sys.stdin:	1✔
73	# Have to close the temp file so that mpy-cross can read it, so we
74	# create our own context manager to delete the file when we are done
75	# using it.
76
77	@contextlib.contextmanager	×
78	def temp_context():	×
79	try:	×
80	with NamedTemporaryFile("wb", suffix=".py", delete=False) as temp:	×
81	temp.write(file.buffer.read())	×
82
83	yield temp.name	×
84	finally:
85	try:	×
86	os.remove(temp.name)	×
87	except NameError:	×
88	# if NamedTemporaryFile() throws, temp is not defined
89	pass	×
90	except OSError:	×
91	# file was already deleted or other strangeness
92	pass	×
93
94	return temp_context()	×
95
96	file.close()	1✔
97	return contextlib.nullcontext(file.name)	1✔
98
99
100	class Compile(Tool):	1✔
101	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
102	parser = subparsers.add_parser(	1✔
103	"compile",
104	help="compile a Pybricks program without running it",
105	)
106	parser.add_argument(	1✔
107	"file",
108	metavar="<file>",
109	help="path to a MicroPython script or `-` for stdin",
110	type=argparse.FileType(encoding="utf-8"),
111	)
112	parser.add_argument(	1✔
113	"--abi",
114	metavar="<n>",
115	help="the MPY ABI version, one of %(choices)s (default: %(default)s)",
116	default=6,
117	choices=[5, 6],
118	type=int,
119	)
120	parser.add_argument(	1✔
121	"--bin",
122	action="store_true",
123	help="output unformatted binary data only (useful for pipes)",
124	)
125	parser.tool = self	1✔
126
127	async def run(self, args: argparse.Namespace):	1✔
128	from pybricksdev.compile import compile_multi_file, print_mpy	1✔
129
130	with _get_script_path(args.file) as script_path:	1✔
131	mpy = await compile_multi_file(script_path, args.abi)	1✔
132	if args.bin:	1✔
133	sys.stdout.buffer.write(mpy)	×
134	return	×
135	print_mpy(mpy)	1✔
136
137
138	class Run(Tool):	1✔
139	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
140	parser = subparsers.add_parser(	1✔
141	"run",
142	help="run a Pybricks program",
143	)
144	parser.tool = self	1✔
145	parser.add_argument(	1✔
146	"conntype",
147	metavar="<connection type>",
148	help="connection type: %(choices)s",
149	choices=["ble", "usb"],
150	)
151	parser.add_argument(	1✔
152	"file",
153	metavar="<file>",
154	help="path to a MicroPython script or `-` for stdin",
155	type=argparse.FileType(encoding="utf-8"),
156	)
157	parser.add_argument(	1✔
158	"-n",
159	"--name",
160	metavar="<name>",
161	required=False,
162	help="Bluetooth device name or Bluetooth address for BLE connection; "
163	"serial port name for USB connection",
164	)
165
166	parser.add_argument(	1✔
167	"--start",
168	help="Start the program immediately after downloading it.",
169	action=argparse.BooleanOptionalAction,
170	default=True,
171	)
172
173	parser.add_argument(	1✔
174	"--wait",
175	help="Wait for the program to complete before disconnecting. Only applies when starting program right away.",
176	action=argparse.BooleanOptionalAction,
177	default=True,
178	)
179
180	parser.add_argument(	1✔
181	"--stay-connected",
182	help="Add a menu option to resend the code with bluetooth instead of disconnecting from the robot after the program ends.",
183	action=argparse.BooleanOptionalAction,
184	default=False,
185	)
186
187	async def stay_connected_menu(self, hub: PybricksHub, args: argparse.Namespace):	1✔
188
189	if args.conntype == "ble":	1✔
190	from pybricksdev.ble import find_device as find_ble	1✔
191	from pybricksdev.connections.pybricks import PybricksHubBLE	1✔
192	else:
193	from usb.core import find as find_usb	×
194
195	from pybricksdev.connections.pybricks import PybricksHubUSB	×
196	from pybricksdev.usb import (	×
197	EV3_USB_PID,
198	LEGO_USB_VID,
199	MINDSTORMS_INVENTOR_USB_PID,
200	NXT_USB_PID,
201	SPIKE_ESSENTIAL_USB_PID,
202	SPIKE_PRIME_USB_PID,
203	)
204
205	def is_pybricks_usb(dev):	×
206	return (	×
207	(dev.idVendor == LEGO_USB_VID)
208	and (
209	dev.idProduct
210	in [
211	NXT_USB_PID,
212	EV3_USB_PID,
213	SPIKE_PRIME_USB_PID,
214	SPIKE_ESSENTIAL_USB_PID,
215	MINDSTORMS_INVENTOR_USB_PID,
216	]
217	)
218	and dev.product.endswith("Pybricks")
219	)
220
221	class ResponseOptions(IntEnum):	1✔
222	RECOMPILE_RUN = 0	1✔
223	RECOMPILE_DOWNLOAD = 1	1✔
224	RUN_STORED = 2	1✔
225	CHANGE_TARGET_FILE = 3	1✔
226	EXIT = 4	1✔
227
228	async def reconnect_hub():	1✔
229	if not await questionary.confirm(	1✔
230	"\nThe hub has been disconnected. Would you like to re-connect?"
231	).ask_async():
232	exit()	×
233
234	if args.conntype == "ble":	1✔
235	print(	1✔
236	f"Searching for {args.name or 'any hub with Pybricks service'}..."
237	)
238	device_or_address = await find_ble(args.name)	1✔
239	hub = PybricksHubBLE(device_or_address)	1✔
240	elif args.conntype == "usb":	×
241	device_or_address = find_usb(custom_match=is_pybricks_usb)	×
242	hub = PybricksHubUSB(device_or_address)	×
243
244	await hub.connect()	1✔
245	# re-enable echoing of the hub's stdout
246	hub._enable_line_handler = True	1✔
247	hub.print_output = True	1✔
248	return hub	1✔
249
250	response_options = [	1✔
251	"Recompile and Run",
252	"Recompile and Download",
253	"Run Stored Program",
254	"Change Target File",
255	"Exit",
256	]
257	# the entry that is selected by default when the menu opens
258	# this is overridden after the user picks an option
259	# so that the default option is always the one that was last chosen
260	default_response_option = (	1✔
261	ResponseOptions.RECOMPILE_RUN
262	if args.start
263	else ResponseOptions.RECOMPILE_DOWNLOAD
264	)
265
266	while True:	1✔
267	try:	1✔
268	if args.file is sys.stdin:	1✔
269	await hub.race_disconnect(	×
270	hub.race_power_button_press(
271	questionary.press_any_key_to_continue(
272	"The hub will stay connected and echo its output to the terminal. Press any key to exit."
273	).ask_async()
274	)
275	)
276	return	×
277	response = await hub.race_disconnect(	1✔
278	hub.race_power_button_press(
279	questionary.select(
280	f"Would you like to re-compile {os.path.basename(args.file.name)}?",
281	response_options,
282	default=(response_options[default_response_option]),
283	).ask_async()
284	)
285	)
286
287	default_response_option = response_options.index(response)	1✔
288
289	match response_options.index(response):	1✔
290
291	case ResponseOptions.RECOMPILE_RUN:	1✔
292	with _get_script_path(args.file) as script_path:	1✔
293	await hub.run(script_path, wait=True)	1✔
294
295	case ResponseOptions.RECOMPILE_DOWNLOAD:	1✔
296	with _get_script_path(args.file) as script_path:	1✔
297	await hub.download(script_path)	1✔
298
299	case ResponseOptions.RUN_STORED:	1✔
300	if hub.fw_version < Version("3.2.0-beta.4"):	1✔
301	print(	1✔
302	"Running a stored program remotely is only supported in the hub firmware version >= v3.2.0."
303	)
304	else:
305	await hub.start_user_program()	1✔
306	await hub._wait_for_user_program_stop()	1✔
307
308	case ResponseOptions.CHANGE_TARGET_FILE:	1✔
309	args.file.close()	1✔
310	while True:	1✔
311	try:	1✔
312	args.file = open(	1✔
313	await hub.race_disconnect(
314	hub.race_power_button_press(
315	questionary.path(
316	"What file would you like to use?"
317	).ask_async()
318	)
319	),
320	encoding="utf-8",
321	)
322	break	1✔
323	except FileNotFoundError:	×
324	print("The file was not found. Please try again.")	×
325	# send the new target file to the hub
326	with _get_script_path(args.file) as script_path:	1✔
327	await hub.download(script_path)	1✔
328
329	case _:	1✔
330	return	1✔
331
332	except subprocess.CalledProcessError as e:	1✔
333	print()	×
334	print("mpy-cross failed to compile the program:")	×
335	print(e.stderr.decode())	×
336
337	except HubPowerButtonPressedError:	1✔
338	# This means the user pressed the button on the hub to re-start the
339	# current program, so the menu was canceled and we are now printing
340	# the hub stdout until the user program ends on the hub.
341	try:	1✔
342	await hub._wait_for_power_button_release()	1✔
343	await hub._wait_for_user_program_stop()	1✔
344
345	except HubDisconnectError:	×
346	hub = await reconnect_hub()	×
347
348	except HubDisconnectError:	1✔
349	# let terminal cool off before making a new prompt
350	await asyncio.sleep(0.3)	1✔
351	hub = await reconnect_hub()	1✔
352
353	async def run(self, args: argparse.Namespace):	1✔
354
355	# Pick the right connection
356	if args.conntype == "ble":	1✔
357	from pybricksdev.ble import find_device as find_ble	1✔
358	from pybricksdev.connections.pybricks import PybricksHubBLE	1✔
359
360	# It is a Pybricks Hub with BLE. Device name or address is given.
361	print(f"Searching for {args.name or 'any hub with Pybricks service'}...")	1✔
362	device_or_address = await find_ble(args.name)	1✔
363	hub = PybricksHubBLE(device_or_address)	1✔
364	elif args.conntype == "usb":	1✔
365	from usb.core import find as find_usb	1✔
366
367	from pybricksdev.connections.pybricks import PybricksHubUSB	1✔
368	from pybricksdev.usb import (	1✔
369	EV3_USB_PID,
370	LEGO_USB_VID,
371	MINDSTORMS_INVENTOR_USB_PID,
372	NXT_USB_PID,
373	SPIKE_ESSENTIAL_USB_PID,
374	SPIKE_PRIME_USB_PID,
375	)
376
377	def is_pybricks_usb(dev):	1✔
378	return (	×
379	(dev.idVendor == LEGO_USB_VID)
380	and (
381	dev.idProduct
382	in [
383	NXT_USB_PID,
384	EV3_USB_PID,
385	SPIKE_PRIME_USB_PID,
386	SPIKE_ESSENTIAL_USB_PID,
387	MINDSTORMS_INVENTOR_USB_PID,
388	]
389	)
390	and dev.product.endswith("Pybricks")
391	)
392
393	device_or_address = find_usb(custom_match=is_pybricks_usb)	1✔
394
395	if device_or_address is None:	1✔
396	print("Pybricks Hub not found.", file=sys.stderr)	×
397	exit(1)	×
398
399	hub = PybricksHubUSB(device_or_address)	1✔
400	else:
401	raise ValueError(f"Unknown connection type: {args.conntype}")	×
402
403	# Connect to the address and run the script
404	await hub.connect()	1✔
405	try:	1✔
406	with _get_script_path(args.file) as script_path:	1✔
407	if args.start:	1✔
408	await hub.run(script_path, args.wait or args.stay_connected)	1✔
409	else:
410	if args.stay_connected:	1✔
411	# if the user later starts the program by pressing the button on the hub,
412	# we still want the hub stdout to print to Python's stdout
413	hub.print_output = True	×
414	hub._enable_line_handler = True	×
415	await hub.download(script_path)	1✔
416
417	except subprocess.CalledProcessError as e:	1✔
418	print()	1✔
419	print("mpy-cross failed to compile the program:")	1✔
420	print(e.stderr.decode())	1✔
421
422	finally:
423	if args.stay_connected:	1✔
424	await self.stay_connected_menu(hub, args)	1✔
425	await hub.disconnect()	1✔
426
427
428	class Flash(Tool):	1✔
429
430	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
431	parser = subparsers.add_parser(	×
432	"flash", help="flash firmware on a LEGO Powered Up device"
433	)
434	parser.tool = self	×
435
436	parser.add_argument(	×
437	"firmware",
438	metavar="<firmware-file>",
439	type=argparse.FileType(mode="rb"),
440	help="the firmware .zip file",
441	).completer = FilesCompleter(allowednames=(".zip",))
442
443	parser.add_argument(	×
444	"-n", "--name", metavar="<name>", type=str, help="a custom name for the hub"
445	)
446
447	def run(self, args: argparse.Namespace):	1✔
448	from pybricksdev.cli.flash import flash_firmware	×
449
450	return flash_firmware(args.firmware, args.name)	×
451
452
453	class DFUBackup(Tool):	1✔
454	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
455	parser = subparsers.add_parser("backup", help="backup firmware using DFU")	×
456	parser.tool = self	×
457	parser.add_argument(	×
458	"firmware",
459	metavar="<firmware-file>",
460	type=argparse.FileType(mode="wb"),
461	help="the firmware .bin file",
462	).completer = FilesCompleter(allowednames=(".bin",))
463
464	async def run(self, args: argparse.Namespace):	1✔
465	from pybricksdev.dfu import backup_dfu	×
466
467	backup_dfu(args.firmware)	×
468
469
470	class DFURestore(Tool):	1✔
471	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
472	parser = subparsers.add_parser(	×
473	"restore",
474	help="restore firmware using DFU",
475	)
476	parser.tool = self	×
477	parser.add_argument(	×
478	"firmware",
479	metavar="<firmware-file>",
480	type=argparse.FileType(mode="rb"),
481	help="the firmware .bin file",
482	).completer = FilesCompleter(allowednames=(".bin",))
483
484	async def run(self, args: argparse.Namespace):	1✔
485	from pybricksdev.dfu import restore_dfu	×
486
487	restore_dfu(args.firmware)	×
488
489
490	class DFU(Tool):	1✔
491	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
492	self.parser = subparsers.add_parser(	×
493	"dfu",
494	help="use DFU to backup or restore firmware",
495	)
496	self.parser.tool = self	×
497	self.subparsers = self.parser.add_subparsers(	×
498	metavar="<action>", dest="action", help="the action to perform"
499	)
500
501	for tool in DFUBackup(), DFURestore():	×
502	tool.add_parser(self.subparsers)	×
503
504	def run(self, args: argparse.Namespace):	1✔
505	if args.action not in self.subparsers.choices:	×
506	self.parser.error(	×
507	f'Missing name of action: {"\|".join(self.subparsers.choices.keys())}'
508	)
509
510	return self.subparsers.choices[args.action].tool.run(args)	×
511
512
513	class OADFlash(Tool):	1✔
514	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
515	parser = subparsers.add_parser(	×
516	"flash",
517	help="update firmware on a LEGO Powered Up device using TI OAD",
518	)
519	parser.tool = self	×
520	parser.add_argument(	×
521	"firmware",
522	metavar="<firmware-file>",
523	type=argparse.FileType(mode="rb"),
524	help="the firmware .oda file",
525	).completer = FilesCompleter(allowednames=(".oda",))
526
527	async def run(self, args: argparse.Namespace):	1✔
528	from pybricksdev.cli.oad import flash_oad_image	×
529
530	await flash_oad_image(args.firmware)	×
531
532
533	class OADInfo(Tool):	1✔
534	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
535	parser = subparsers.add_parser(	×
536	"info",
537	help="get information about firmware on a LEGO Powered Up device using TI OAD",
538	)
539	parser.tool = self	×
540
541	async def run(self, args: argparse.Namespace):	1✔
542	from pybricksdev.cli.oad import dump_oad_info	×
543
544	await dump_oad_info()	×
545
546
547	class OAD(Tool):	1✔
548	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
549	self.parser = subparsers.add_parser(	×
550	"oad",
551	help="update firmware on a LEGO Powered Up device using TI OAD",
552	)
553	self.parser.tool = self	×
554	self.subparsers = self.parser.add_subparsers(	×
555	metavar="<action>", dest="action", help="the action to perform"
556	)
557
558	for tool in OADFlash(), OADInfo():	×
559	tool.add_parser(self.subparsers)	×
560
561	def run(self, args: argparse.Namespace):	1✔
562	if args.action not in self.subparsers.choices:	×
563	self.parser.error(	×
564	f'Missing name of action: {"\|".join(self.subparsers.choices.keys())}'
565	)
566
567	return self.subparsers.choices[args.action].tool.run(args)	×
568
569
570	class LWP3Repl(Tool):	1✔
571	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
572	parser = subparsers.add_parser(	×
573	"repl",
574	help="interactive REPL for sending and receiving LWP3 messages",
575	)
576	parser.tool = self	×
577
578	def run(self, args: argparse.Namespace):	1✔
579	from pybricksdev.cli.lwp3.repl import repl, setup_repl_logging	×
580
581	setup_repl_logging()	×
582	return repl()	×
583
584
585	class LWP3(Tool):	1✔
586	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
587	self.parser = subparsers.add_parser(	×
588	"lwp3", help="interact with devices using LWP3"
589	)
590	self.parser.tool = self	×
591	self.subparsers = self.parser.add_subparsers(	×
592	metavar="<lwp3-tool>", dest="lwp3_tool", help="the tool to run"
593	)
594
595	for tool in (LWP3Repl(),):	×
596	tool.add_parser(self.subparsers)	×
597
598	def run(self, args: argparse.Namespace):	1✔
599	if args.lwp3_tool not in self.subparsers.choices:	×
600	self.parser.error(	×
601	f'Missing name of tool: {"\|".join(self.subparsers.choices.keys())}'
602	)
603
604	return self.subparsers.choices[args.lwp3_tool].tool.run(args)	×
605
606
607	class Udev(Tool):	1✔
608	def add_parser(self, subparsers: argparse._SubParsersAction):	1✔
609	parser = subparsers.add_parser("udev", help="print udev rules to stdout")	1✔
610	parser.tool = self	1✔
611
612	async def run(self, args: argparse.Namespace):	1✔
613	from importlib.resources import read_text	1✔
614
615	from pybricksdev import resources	1✔
616
617	print(read_text(resources, resources.UDEV_RULES))	1✔
618
619
620	def main():	1✔
621	"""Runs ``pybricksdev`` command line interface."""
622
623	if sys.platform == "win32":	×
624	# Hack around bad side-effects of pythoncom on Windows
625	try:	×
626	from bleak_winrt._winrt import MTA, init_apartment	×
627	except ImportError:	×
628	from winrt._winrt import MTA, init_apartment	×
629
630	init_apartment(MTA)	×
631
632	# Provide main description and help.
633	parser = argparse.ArgumentParser(	×
634	prog=PROG_NAME,
635	description="Utilities for Pybricks developers.",
636	epilog="Run `%(prog)s <tool> --help` for tool-specific arguments.",
637	)
638
639	parser.add_argument(	×
640	"-v", "--version", action="version", version=f"{MODULE_NAME} v{MODULE_VERSION}"
641	)
642	parser.add_argument(	×
643	"-d", "--debug", action="store_true", help="enable debug logging"
644	)
645
646	subparsers = parser.add_subparsers(	×
647	metavar="<tool>",
648	dest="tool",
649	help="the tool to use",
650	)
651
652	for tool in Compile(), Run(), Flash(), DFU(), OAD(), LWP3(), Udev():	×
653	tool.add_parser(subparsers)	×
654
655	argcomplete.autocomplete(parser)	×
656	args = parser.parse_args()	×
657
658	logging.basicConfig(	×
659	format="%(asctime)s: %(levelname)s: %(name)s: %(message)s",
660	level=logging.DEBUG if args.debug else logging.WARNING,
661	)
662
663	if not args.tool:	×
664	parser.error(f'Missing name of tool: {"\|".join(subparsers.choices.keys())}')	×
665
666	asyncio.run(subparsers.choices[args.tool].tool.run(args))	×

pybricks / pybricksdev / 20311043488

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous