From 985b3d9e60b26f971b94890e752b90c94bb216ae Mon Sep 17 00:00:00 2001 From: Nate Karstens Date: Sat, 25 Nov 2023 08:00:00 -0600 Subject: [PATCH 1/3] Move PybricksHub client creation into constructor Moves client creation into the PybricksHub constructor. A future change will move this into a subclass specific for BLE or USB, so configuring this is more appropriately done in the constructor instead of the connect method. The disconnect logic is moved into a class method so that it can be shared by both BLE and USB implementations. Temporarily comments out a message that prints the device name as this is no longer available in the connect method. Signed-off-by: Nate Karstens --- pybricksdev/cli/__init__.py | 7 +++---- pybricksdev/connections/ev3dev.py | 9 ++++++--- pybricksdev/connections/pybricks.py | 25 +++++++++++++------------ 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/pybricksdev/cli/__init__.py b/pybricksdev/cli/__init__.py index 8fe8fe3..6cd1cac 100644 --- a/pybricksdev/cli/__init__.py +++ b/pybricksdev/cli/__init__.py @@ -183,22 +183,21 @@ async def run(self, args: argparse.Namespace): print("--name is required for SSH connections", file=sys.stderr) exit(1) - hub = EV3Connection() device_or_address = socket.gethostbyname(args.name) + hub = EV3Connection(device_or_address) elif args.conntype == "ble": # It is a Pybricks Hub with BLE. Device name or address is given. - hub = PybricksHub() print(f"Searching for {args.name or 'any hub with Pybricks service'}...") device_or_address = await find_device(args.name) + hub = PybricksHub(device_or_address) elif args.conntype == "usb": hub = REPLHub() - device_or_address = None else: raise ValueError(f"Unknown connection type: {args.conntype}") # Connect to the address and run the script - await hub.connect(device_or_address) + await hub.connect() try: with _get_script_path(args.file) as script_path: await hub.run(script_path, args.wait) diff --git a/pybricksdev/connections/ev3dev.py b/pybricksdev/connections/ev3dev.py index ee9b0aa..084ce06 100644 --- a/pybricksdev/connections/ev3dev.py +++ b/pybricksdev/connections/ev3dev.py @@ -18,10 +18,13 @@ class EV3Connection: _USER = "robot" _PASSWORD = "maker" + def __init__(self, address: str): + self._address = address + def abs_path(self, path): return os.path.join(self._HOME, path) - async def connect(self, address): + async def connect(self): """Connects to ev3dev using SSH with a known IP address. Arguments: @@ -33,9 +36,9 @@ async def connect(self, address): Connect failed. """ - print("Connecting to", address, "...", end=" ") + print("Connecting to", self._address, "...", end=" ") self.client = await asyncssh.connect( - address, username=self._USER, password=self._PASSWORD + self._address, username=self._USER, password=self._PASSWORD ) print("Connected.", end=" ") self.client.sftp = await self.client.start_sftp_client() diff --git a/pybricksdev/connections/pybricks.py b/pybricksdev/connections/pybricks.py index 0f61a88..c43b3da 100644 --- a/pybricksdev/connections/pybricks.py +++ b/pybricksdev/connections/pybricks.py @@ -78,7 +78,7 @@ class PybricksHub: has not been connected yet or the connected hub has Pybricks profile < v1.2.0. """ - def __init__(self): + def __init__(self, device: BLEDevice): self.connection_state_observable = BehaviorSubject(ConnectionState.DISCONNECTED) self.status_observable = BehaviorSubject(StatusFlag(0)) self._stdout_subject = Subject() @@ -120,6 +120,11 @@ def __init__(self): # File handle for logging self.log_file = None + def handle_disconnect(_: BleakClient): + self._handle_disconnect() + + self.client = BleakClient(device, disconnected_callback=handle_disconnect) + @property def stdout_observable(self) -> Observable[bytes]: """ @@ -227,18 +232,20 @@ def _pybricks_service_handler(self, _: int, data: bytes) -> None: if self._enable_line_handler: self._handle_line_data(payload) - async def connect(self, device: BLEDevice): - """Connects to a device that was discovered with :meth:`pybricksdev.ble.find_device` + def _handle_disconnect(self): + logger.info("Disconnected!") + self.connection_state_observable.on_next(ConnectionState.DISCONNECTED) - Args: - device: The device to connect to. + async def connect(self): + """Connects to a device that was discovered with :meth:`pybricksdev.ble.find_device` Raises: BleakError: if connecting failed (or old firmware without Device Information Service) RuntimeError: if Pybricks Protocol version is not supported """ - logger.info(f"Connecting to {device.name}") + # TODO: Fix this + # logger.info(f"Connecting to {device.name}") if self.connection_state_observable.value != ConnectionState.DISCONNECTED: raise RuntimeError( @@ -252,12 +259,6 @@ async def connect(self, device: BLEDevice): self.connection_state_observable.on_next, ConnectionState.DISCONNECTED ) - def handle_disconnect(_: BleakClient): - logger.info("Disconnected!") - self.connection_state_observable.on_next(ConnectionState.DISCONNECTED) - - self.client = BleakClient(device, disconnected_callback=handle_disconnect) - await self.client.connect() stack.push_async_callback(self.disconnect) From e82a7f07a9a00f9cf5b33e750185cd3bc799cd67 Mon Sep 17 00:00:00 2001 From: Nate Karstens Date: Sat, 25 Nov 2023 08:00:00 -0600 Subject: [PATCH 2/3] Move BLE logic into new class Moves all logic specific to BLE connections into a new subclass of PybricksHub. Another subclass will be added later to handle USB connections. Code to retrieve firmware version, hub capabilities, etc. is moved into the connect step to better abstract this for any connection medium. Signed-off-by: Nate Karstens --- pybricksdev/cli/__init__.py | 4 +- pybricksdev/connections/pybricks.py | 149 ++++++++++++++++------------ 2 files changed, 89 insertions(+), 64 deletions(-) diff --git a/pybricksdev/cli/__init__.py b/pybricksdev/cli/__init__.py index 6cd1cac..537e338 100644 --- a/pybricksdev/cli/__init__.py +++ b/pybricksdev/cli/__init__.py @@ -174,7 +174,7 @@ async def run(self, args: argparse.Namespace): from pybricksdev.ble import find_device from pybricksdev.connections.ev3dev import EV3Connection from pybricksdev.connections.lego import REPLHub - from pybricksdev.connections.pybricks import PybricksHub + from pybricksdev.connections.pybricks import PybricksHubBLE # Pick the right connection if args.conntype == "ssh": @@ -189,7 +189,7 @@ async def run(self, args: argparse.Namespace): # 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_device(args.name) - hub = PybricksHub(device_or_address) + hub = PybricksHubBLE(device_or_address) elif args.conntype == "usb": hub = REPLHub() diff --git a/pybricksdev/connections/pybricks.py b/pybricksdev/connections/pybricks.py index c43b3da..a5024f6 100644 --- a/pybricksdev/connections/pybricks.py +++ b/pybricksdev/connections/pybricks.py @@ -6,7 +6,7 @@ import logging import os import struct -from typing import Awaitable, List, Optional, TypeVar +from typing import Awaitable, Callable, List, Optional, TypeVar import reactivex.operators as op import semver @@ -78,7 +78,7 @@ class PybricksHub: has not been connected yet or the connected hub has Pybricks profile < v1.2.0. """ - def __init__(self, device: BLEDevice): + def __init__(self): self.connection_state_observable = BehaviorSubject(ConnectionState.DISCONNECTED) self.status_observable = BehaviorSubject(StatusFlag(0)) self._stdout_subject = Subject() @@ -120,11 +120,6 @@ def __init__(self, device: BLEDevice): # File handle for logging self.log_file = None - def handle_disconnect(_: BleakClient): - self._handle_disconnect() - - self.client = BleakClient(device, disconnected_callback=handle_disconnect) - @property def stdout_observable(self) -> Observable[bytes]: """ @@ -237,16 +232,6 @@ def _handle_disconnect(self): self.connection_state_observable.on_next(ConnectionState.DISCONNECTED) async def connect(self): - """Connects to a device that was discovered with :meth:`pybricksdev.ble.find_device` - - Raises: - BleakError: if connecting failed (or old firmware without Device - Information Service) - RuntimeError: if Pybricks Protocol version is not supported - """ - # TODO: Fix this - # logger.info(f"Connecting to {device.name}") - if self.connection_state_observable.value != ConnectionState.DISCONNECTED: raise RuntimeError( f"attempting to connect with invalid state: {self.connection_state_observable.value}" @@ -259,48 +244,12 @@ async def connect(self): self.connection_state_observable.on_next, ConnectionState.DISCONNECTED ) - await self.client.connect() + await self._client_connect() stack.push_async_callback(self.disconnect) - logger.info("Connected successfully!") - - fw_version = await self.client.read_gatt_char(FW_REV_UUID) - self.fw_version = Version(fw_version.decode()) - - protocol_version = await self.client.read_gatt_char(SW_REV_UUID) - protocol_version = semver.VersionInfo.parse(protocol_version.decode()) - - if ( - protocol_version < PYBRICKS_PROTOCOL_VERSION - or protocol_version >= PYBRICKS_PROTOCOL_VERSION.bump_major() - ): - raise RuntimeError( - f"Unsupported Pybricks protocol version: {protocol_version}" - ) - - pnp_id = await self.client.read_gatt_char(PNP_ID_UUID) - _, _, self.hub_kind, self.hub_variant = unpack_pnp_id(pnp_id) - - if protocol_version >= "1.2.0": - caps = await self.client.read_gatt_char(PYBRICKS_HUB_CAPABILITIES_UUID) - ( - self._max_write_size, - self._capability_flags, - self._max_user_program_size, - ) = unpack_hub_capabilities(caps) - else: - # HACK: prior to profile v1.2.0 isn't a proper way to get the - # MPY ABI version from hub so we use heuristics on the firmware version - self._mpy_abi_version = ( - 6 if self.fw_version >= Version("3.2.0b2") else 5 - ) - - if protocol_version < "1.3.0": - self._legacy_stdio = True - - await self.client.start_notify(NUS_TX_UUID, self._nus_handler) - await self.client.start_notify( + await self.start_notify(NUS_TX_UUID, self._nus_handler) + await self.start_notify( PYBRICKS_COMMAND_EVENT_UUID, self._pybricks_service_handler ) @@ -314,7 +263,7 @@ async def disconnect(self): if self.connection_state_observable.value == ConnectionState.CONNECTED: self.connection_state_observable.on_next(ConnectionState.DISCONNECTING) - await self.client.disconnect() + await self._client_disconnect() # ConnectionState.DISCONNECTED should be set by disconnect callback assert ( self.connection_state_observable.value == ConnectionState.DISCONNECTED @@ -453,7 +402,7 @@ async def download_user_program(self, program: bytes) -> None: ) # clear user program meta so hub doesn't try to run invalid program - await self.client.write_gatt_char( + await self.write_gatt_char( PYBRICKS_COMMAND_EVENT_UUID, struct.pack(" None: total=len(program), unit="B", unit_scale=True ) as pbar: for i, c in enumerate(chunk(program, payload_size)): - await self.client.write_gatt_char( + await self.write_gatt_char( PYBRICKS_COMMAND_EVENT_UUID, struct.pack( f" None: pbar.update(len(c)) # set the metadata to notify that writing was successful - await self.client.write_gatt_char( + await self.write_gatt_char( PYBRICKS_COMMAND_EVENT_UUID, struct.pack(" None: Requires hub with Pybricks Profile >= v1.2.0. """ - await self.client.write_gatt_char( + await self.write_gatt_char( PYBRICKS_COMMAND_EVENT_UUID, struct.pack(" None: """ Stops the user program on the hub if it is running. """ - await self.client.write_gatt_char( + await self.write_gatt_char( PYBRICKS_COMMAND_EVENT_UUID, struct.pack(" bool: + """Connects to a device that was discovered with :meth:`pybricksdev.ble.find_device` + + Raises: + BleakError: if connecting failed (or old firmware without Device + Information Service) + RuntimeError: if Pybricks Protocol version is not supported + """ + + logger.info(f"Connecting to {self._device.name}") + await self._client.connect() + logger.info("Connected successfully!") + + fw_version = await self.read_gatt_char(FW_REV_UUID) + self.fw_version = Version(fw_version.decode()) + + protocol_version = await self.read_gatt_char(SW_REV_UUID) + protocol_version = semver.VersionInfo.parse(protocol_version.decode()) + + if ( + protocol_version < PYBRICKS_PROTOCOL_VERSION + or protocol_version >= PYBRICKS_PROTOCOL_VERSION.bump_major() + ): + raise RuntimeError( + f"Unsupported Pybricks protocol version: {protocol_version}" + ) + + pnp_id = await self.read_gatt_char(PNP_ID_UUID) + _, _, self.hub_kind, self.hub_variant = unpack_pnp_id(pnp_id) + + if protocol_version >= "1.2.0": + caps = await self.read_gatt_char(PYBRICKS_HUB_CAPABILITIES_UUID) + ( + self._max_write_size, + self._capability_flags, + self._max_user_program_size, + ) = unpack_hub_capabilities(caps) + else: + # HACK: prior to profile v1.2.0 isn't a proper way to get the + # MPY ABI version from hub so we use heuristics on the firmware version + self._mpy_abi_version = 6 if self.fw_version >= Version("3.2.0b2") else 5 + + if protocol_version < "1.3.0": + self._legacy_stdio = True + + return True + + async def _client_disconnect(self) -> bool: + return await self._client.disconnect() + + async def read_gatt_char(self, uuid: str) -> bytearray: + return await self._client.read_gatt_char(uuid) + + async def write_gatt_char(self, uuid: str, data, response: bool) -> None: + return await self._client.write_gatt_char(uuid, data, response) + + async def start_notify(self, uuid: str, callback: Callable) -> None: + return await self._client.start_notify(uuid, callback) From 9b5739bc1d275ef8b302b24066627af23a702aa9 Mon Sep 17 00:00:00 2001 From: Nate Karstens Date: Sat, 25 Nov 2023 08:00:00 -0600 Subject: [PATCH 3/3] Add support for USB connections Adds a new subclass of PybricksHub that manages USB connections. Co-developed-by: David Lechner Signed-off-by: Nate Karstens --- CHANGELOG.md | 3 + pybricksdev/cli/__init__.py | 45 +++++++-- pybricksdev/connections/pybricks.py | 148 ++++++++++++++++++++++++++++ pybricksdev/usb/__init__.py | 3 + pybricksdev/usb/pybricks.py | 30 ++++++ 5 files changed, 222 insertions(+), 7 deletions(-) create mode 100644 pybricksdev/usb/pybricks.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 52814f5..8c03cde 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- Partial/experimental support for `pybricksdev run usb`. + ### Fixed - Fix crash when running `pybricksdev run ble -` (bug introduced in alpha.49). diff --git a/pybricksdev/cli/__init__.py b/pybricksdev/cli/__init__.py index 537e338..4e2cdcb 100644 --- a/pybricksdev/cli/__init__.py +++ b/pybricksdev/cli/__init__.py @@ -171,13 +171,11 @@ def add_parser(self, subparsers: argparse._SubParsersAction): ) async def run(self, args: argparse.Namespace): - from pybricksdev.ble import find_device - from pybricksdev.connections.ev3dev import EV3Connection - from pybricksdev.connections.lego import REPLHub - from pybricksdev.connections.pybricks import PybricksHubBLE # Pick the right connection if args.conntype == "ssh": + from pybricksdev.connections.ev3dev import EV3Connection + # So it's an ev3dev if args.name is None: print("--name is required for SSH connections", file=sys.stderr) @@ -186,13 +184,46 @@ async def run(self, args: argparse.Namespace): device_or_address = socket.gethostbyname(args.name) hub = EV3Connection(device_or_address) elif 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_device(args.name) + device_or_address = await find_ble(args.name) hub = PybricksHubBLE(device_or_address) - elif args.conntype == "usb": - hub = REPLHub() + from usb.core import find as find_usb + + from pybricksdev.connections.pybricks import PybricksHubUSB + from pybricksdev.usb import ( + LEGO_USB_VID, + MINDSTORMS_INVENTOR_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 [ + 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 not None: + hub = PybricksHubUSB(device_or_address) + else: + from pybricksdev.connections.lego import REPLHub + + hub = REPLHub() else: raise ValueError(f"Unknown connection type: {args.conntype}") diff --git a/pybricksdev/connections/pybricks.py b/pybricksdev/connections/pybricks.py index a5024f6..f998ed0 100644 --- a/pybricksdev/connections/pybricks.py +++ b/pybricksdev/connections/pybricks.py @@ -7,6 +7,7 @@ import os import struct from typing import Awaitable, Callable, List, Optional, TypeVar +from uuid import UUID import reactivex.operators as op import semver @@ -17,6 +18,10 @@ from reactivex.subject import BehaviorSubject, Subject from tqdm.auto import tqdm from tqdm.contrib.logging import logging_redirect_tqdm +from usb.control import get_descriptor +from usb.core import Device as USBDevice +from usb.core import Endpoint, USBTimeoutError +from usb.util import ENDPOINT_IN, ENDPOINT_OUT, endpoint_direction, find_descriptor from pybricksdev.ble.lwp3.bytecodes import HubKind from pybricksdev.ble.nus import NUS_RX_UUID, NUS_TX_UUID @@ -38,6 +43,10 @@ from pybricksdev.connections import ConnectionState from pybricksdev.tools import chunk from pybricksdev.tools.checksum import xor_bytes +from pybricksdev.usb.pybricks import ( + PybricksUsbInEpMessageType, + PybricksUsbOutEpMessageType, +) logger = logging.getLogger(__name__) @@ -707,3 +716,142 @@ async def write_gatt_char(self, uuid: str, data, response: bool) -> None: async def start_notify(self, uuid: str, callback: Callable) -> None: return await self._client.start_notify(uuid, callback) + + +class PybricksHubUSB(PybricksHub): + _device: USBDevice + _ep_in: Endpoint + _ep_out: Endpoint + _notify_callbacks: dict[str, Callable] = {} + _monitor_task: asyncio.Task + + def __init__(self, device: USBDevice): + super().__init__() + self._device = device + self._response_queue = asyncio.Queue[bytes]() + + async def _client_connect(self) -> bool: + # Reset is essential to ensure endpoints are in a good state. + self._device.reset() + self._device.set_configuration() + + # Save input and output endpoints + cfg = self._device.get_active_configuration() + intf = cfg[(0, 0)] + self._ep_in = find_descriptor( + intf, + custom_match=lambda e: endpoint_direction(e.bEndpointAddress) + == ENDPOINT_IN, + ) + self._ep_out = find_descriptor( + intf, + custom_match=lambda e: endpoint_direction(e.bEndpointAddress) + == ENDPOINT_OUT, + ) + + # There is 1 byte overhead for PybricksUsbMessageType + self._max_write_size = self._ep_out.wMaxPacketSize - 1 + + # Get length of BOS descriptor + bos_descriptor = get_descriptor(self._device, 5, 0x0F, 0) + (ofst, _, bos_len, _) = struct.unpack(" bool: + self._monitor_task.cancel() + self._handle_disconnect() + + async def read_gatt_char(self, uuid: str) -> bytearray: + # Most stuff is available via other properties due to reading BOS + # descriptor during connect. + raise NotImplementedError + + async def write_gatt_char(self, uuid: str, data, response: bool) -> None: + if uuid.lower() != PYBRICKS_COMMAND_EVENT_UUID: + raise ValueError("Only Pybricks command event UUID is supported") + + if not response: + raise ValueError("Response is required for USB") + + self._ep_out.write(bytes([PybricksUsbOutEpMessageType.COMMAND]) + data) + # FIXME: This needs to race with hub disconnect, and could also use a + # timeout, otherwise it blocks forever. Pyusb doesn't currently seem to + # have any disconnect callback. + reply = await self._response_queue.get() + + # REVISIT: could look up status error code and convert to string, + # although BLE doesn't do that either. + if int.from_bytes(reply[:4], "little") != 0: + raise RuntimeError(f"Write failed: {reply[0]}") + + async def start_notify(self, uuid: str, callback: Callable) -> None: + # TODO: need to send subscribe message over USB + self._notify_callbacks[uuid] = callback + + async def _monitor_usb(self): + loop = asyncio.get_running_loop() + + while True: + msg = await loop.run_in_executor(None, self._read_usb) + + if msg is None: + continue + + if len(msg) == 0: + logger.warning("Empty USB message") + continue + + if msg[0] == PybricksUsbInEpMessageType.RESPONSE: + self._response_queue.put_nowait(msg[1:]) + elif msg[0] == PybricksUsbInEpMessageType.EVENT: + callback = self._notify_callbacks.get(PYBRICKS_COMMAND_EVENT_UUID) + if callback: + callback(None, msg[1:]) + else: + logger.warning("Unknown USB message type: %d", msg[0]) + + def _read_usb(self) -> bytes | None: + try: + msg = self._ep_in.read(self._ep_in.wMaxPacketSize) + return msg + except USBTimeoutError: + return None diff --git a/pybricksdev/usb/__init__.py b/pybricksdev/usb/__init__.py index ccce97e..55cfd47 100644 --- a/pybricksdev/usb/__init__.py +++ b/pybricksdev/usb/__init__.py @@ -6,5 +6,8 @@ EV3_USB_PID = 0x0005 EV3_BOOTLOADER_USB_PID = 0x0006 SPIKE_PRIME_DFU_USB_PID = 0x0008 +SPIKE_PRIME_USB_PID = 0x0009 SPIKE_ESSENTIAL_DFU_USB_PID = 0x000C +SPIKE_ESSENTIAL_USB_PID = 0x000D +MINDSTORMS_INVENTOR_USB_PID = 0x0010 MINDSTORMS_INVENTOR_DFU_USB_PID = 0x0011 diff --git a/pybricksdev/usb/pybricks.py b/pybricksdev/usb/pybricks.py new file mode 100644 index 0000000..bfbdbf7 --- /dev/null +++ b/pybricksdev/usb/pybricks.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025 The Pybricks Authors + +""" +Pybricks-specific USB protocol. +""" + +from enum import IntEnum + + +class PybricksUsbInEpMessageType(IntEnum): + RESPONSE = 1 + """ + Analogous to BLE status response. + """ + EVENT = 2 + """ + Analogous to BLE notification. + """ + + +class PybricksUsbOutEpMessageType(IntEnum): + SUBSCRIBE = 1 + """ + Analogous to BLE Client Characteristic Configuration Descriptor (CCCD). + """ + COMMAND = 2 + """ + Analogous to BLE write without response. + """