pybricks.compile: add support for MPY v6.1

MicroPython v1.20 bumped the MPY version to v6.1 which only affects
native modules. This add support for a new flag from the hub to
indicate native module support and the new MPY version.

Author: dlech

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Added `PybricksHub.download_user_program()` method ([support#284]).
 - Added support for flashing EV3 firmware.
+- Added support for MPY ABI v6.1.
 
 [support#284]: https://github.com/pybricks/support/issues/284
 

pybricksdev/ble/pybricks.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2021 The Pybricks Authors
+# Copyright (c) 2021-2023 The Pybricks Authors
 
 """
 This module is used for Bluetooth Low Energy communications with devices
@@ -272,9 +272,19 @@ class HubCapabilityFlag(IntFlag):
     """
     Hub supports user programs using a multi-file blob with MicroPython MPY (ABI V6) files.
 
+    Native module support is not implied by this flag but may be indicated by additional flags.
+
     .. availability:: Since Pybricks protocol v1.2.0.
     """
 
+    USER_PROG_MULTI_FILE_MPY6_1_NATIVE = 2 << 1
+    """
+    Hub supports user programs using a multi-file blob with MicroPython MPY (ABI V6.1) files
+    including native module support.
+
+    .. availability:: Since Pybricks protocol v1.3.0.
+    """
+
 
 def unpack_hub_capabilities(data: bytes) -> Tuple[int, HubCapabilityFlag, int]:
     """

pybricksdev/compile.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2019-2022 The Pybricks Authors
+# Copyright (c) 2019-2023 The Pybricks Authors
 
 import asyncio
 import logging
 import os
 from modulefinder import ModuleFinder
-from typing import List, Optional
+from typing import List, Optional, Tuple, Union
 
 import mpy_cross_v5
 import mpy_cross_v6
@@ -76,7 +76,7 @@ async def compile_file(path: str, abi: int, compile_args: Optional[List[str]] =
         return mpy
 
 
-async def compile_multi_file(path: str, abi: int):
+async def compile_multi_file(path: str, abi: Union[int, Tuple[int, int]]):
     """Compiles a Python file and its dependencies with ``mpy-cross``.
 
     On the hub, all dependencies behave as independent modules. Any (leading)
@@ -101,7 +101,8 @@ async def compile_multi_file(path: str, abi: int):
         path:
             Path to script that is to be compiled.
         abi:
-            Expected MPY ABI version.
+            Expected MPY ABI version. Can be major version (int) if no native
+            .mpy modules or tuple of major, minor version.
 
     Returns:
         Concatenation of all compiled files in the format given above.
@@ -123,10 +124,12 @@ async def compile_multi_file(path: str, abi: int):
     # Get a data blob with all scripts.
     parts: List[bytes] = []
 
+    abi_major, abi_minor = (abi, None) if isinstance(abi, int) else abi
+
     for name, module in finder.modules.items():
         if not module.__file__:
             continue
-        mpy = await compile_file(module.__file__, abi)
+        mpy = await compile_file(module.__file__, abi_major)
 
         parts.append(len(mpy).to_bytes(4, "little"))
         parts.append(name.encode() + b"\x00")
@@ -138,9 +141,19 @@ async def compile_multi_file(path: str, abi: int):
             try:
                 with open(os.path.join(spath, f"{name}.mpy"), "rb") as f:
                     mpy = f.read()
-                if mpy[1] != abi:
+                if mpy[1] != abi_major:
+                    raise ValueError(
+                        f"{name}.mpy has abi major version {mpy[1]} while {abi_major} is required"
+                    )
+                if (
+                    abi_minor is not None  # hub indicated a minor version
+                    and (mpy[2] >> 2)  # mpy has native arch
+                    # TODO: How to validate native arch? For now just gets runtime error on hub.
+                    and (mpy[2] & 0x3)
+                    != abi_minor  # mpy minor version does not match hub minor version
+                ):
                     raise ValueError(
-                        f"{name}.mpy has abi version {mpy[1]} while {abi} is required"
+                        f"{name}.mpy has abi minor version {mpy[2] & 0x3} while {abi_minor} is required"
                     )
                 parts.append(len(mpy).to_bytes(4, "little"))
                 parts.append(name.encode() + b"\x00")

pybricksdev/connections/pybricks.py

@@ -548,13 +548,28 @@ async def run(
             return
 
         # since Pybricks profile v1.2.0, the hub will tell us which file format(s) it supports
-        if not (self._capability_flags & HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6):
+        if not (
+            self._capability_flags
+            & (
+                HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6
+                | HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6_1_NATIVE
+            )
+        ):
             raise RuntimeError(
                 "Hub is not compatible with any of the supported file formats"
             )
 
+        # no support for native modules unless one of the flags below is set
+        abi = 6
+
+        if (
+            self._capability_flags
+            & HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6_1_NATIVE
+        ):
+            abi = (6, 1)
+
         if py_path is not None:
-            mpy = await compile_multi_file(py_path, 6)
+            mpy = await compile_multi_file(py_path, abi)
             await self.download_user_program(mpy)
 
         await self.start_user_program()