flash: add option for setting hub name

This adds a new option to set the hub name when flashing firmware. This
requires Pybricks firmware v3.1.0a1 or later.

Author: dlech

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Support for Python 3.9.
 - Short `-n` option for `--name` option in `pybricksdev run`.
+- Option to set hub name when flashing firmware.
 
 ### Changed
 - Update to Bleak v0.12.0.

pybricksdev/cli/__init__.py

@@ -192,18 +192,23 @@ def add_parser(self, subparsers: argparse._SubParsersAction):
             "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"
+        )
+
     async def run(self, args: argparse.Namespace):
         from ..flash import create_firmware
 
         print("Creating firmware")
-        firmware, metadata = await create_firmware(args.firmware)
+        firmware, metadata = await create_firmware(args.firmware, args.name)
 
         if metadata["device-id"] == HubKind.TECHNIC_LARGE:
             from ..dfu import flash_dfu

pybricksdev/flash.py

@@ -11,8 +11,9 @@
 import sys
 import zipfile
 from collections import namedtuple
-from typing import BinaryIO, Dict, Tuple, Union
+from typing import BinaryIO, Dict, Optional, Tuple, Union
 
+import semver
 from tqdm.auto import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
 
@@ -26,17 +27,25 @@
 
 
 async def create_firmware(
-    firmware_zip: Union[str, os.PathLike, BinaryIO]
+    firmware_zip: Union[str, os.PathLike, BinaryIO], name: Optional[str] = None
 ) -> Tuple[bytes, dict]:
     """Creates a firmware blob from base firmware and main.mpy file.
 
     Arguments:
         firmware_zip:
             Path to the firmware zip file or a file-like object.
+        name:
+            A custom name for the hub.
 
     Returns:
         bytes: Composite binary blob with correct padding and checksum.
         dict: Meta data for this firmware file.
+
+    Raises:
+        ValueError:
+            A name is given but the firmware does not support it or the name
+            exceeds the alloted space in the firmware.
+
     """
 
     archive = zipfile.ZipFile(firmware_zip)
@@ -61,6 +70,23 @@ async def create_firmware(
     # pad with 0s to align to 4-byte boundary
     firmware.extend(0 for _ in range(-len(firmware) % 4))
 
+    if name:
+        if semver.compare(metadata["metadata-version"], "1.1.0") < 0:
+            raise ValueError(
+                "this firmware image does not support setting the hub name"
+            )
+
+        name = name.encode() + b"\0"
+
+        max_size = metadata["max-hub-name-size"]
+        if len(name) > max_size:
+            raise ValueError(
+                f"name is too big - must be < {metadata['max-hub-name-size']} UTF-8 bytes"
+            )
+
+        offset = metadata["hub-name-offset"]
+        firmware[offset : offset + len(name)] = name
+
     # append 32-bit little-endian checksum
     if metadata["checksum-type"] == "sum":
         firmware.extend(