compile: Experimental support for multi-file projects.

This makes use of recent firmware changes to allow the main script to
import modules in the same folder. The modules can also import other
modules.

To ensure a smooth transition until a firmware with this feature is
released, scripts without local imports will be downloaded in the old
format (a single raw mpy data blob).

Author: laurensvalk

.vscode/launch.json

@@ -13,7 +13,7 @@
             "args": [
                 "run",
                 "ble",
-                "${workspaceFolder}/demo/shortdemo.py"
+                "${workspaceFolder}/demo/multidemo.py"
             ]
         },
         {

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Added `fw_version` attribute to `pybricksdev.connections.pybricks.PybricksHub`.
+- Experimental support for multi-file projects.
 
 ### Fixed
 - Fixed running programs on hubs with firmware with MPY ABI v5.

demo/module1.py

@@ -0,0 +1,8 @@
+from pybricks.parameters import Color
+
+nice_color = Color.GREEN
+
+print("I am", __name__)
+
+if __name__ == "__main__":
+    print("module1 is main")

demo/module2.py

@@ -0,0 +1,7 @@
+from module1 import nice_color
+
+print("I am", __name__)
+print(nice_color)
+
+if __name__ == "__main__":
+    print("module 2 is main")

demo/multidemo.py

@@ -0,0 +1,11 @@
+from pybricks.hubs import ThisHub
+from pybricks.parameters import Color
+from pybricks.tools import wait
+from module2 import nice_color
+
+hub = ThisHub()
+hub.light.on(Color.RED)
+wait(2000)
+print("I am", __name__)
+hub.light.on(nice_color)
+wait(2000)

pybricksdev/compile.py

@@ -4,6 +4,9 @@
 import asyncio
 import logging
 import os
+import re
+
+from pathlib import Path
 from typing import List, Optional
 
 import mpy_cross_v5
@@ -75,6 +78,108 @@ async def compile_file(path: str, abi: int, compile_args: Optional[List[str]] =
         return mpy
 
 
+async def compile_multi_file(
+    path: str, abi: int, compile_args: Optional[List[str]] = None
+):
+    """Compiles a Python file and its dependencies with ``mpy-cross``.
+
+    All dependencies must be Python modules from the same folder.
+
+    The returned bytes format is of the form:
+
+       - first script size (uint32 little endian)
+       - first script name (no extension, zero terminated string)
+       - first script mpy data
+       - second script size
+       - second script name
+       - second script mpy data
+       - ...
+
+    All components are 0-padded to a size with multiple of 4.
+
+    If the main script does not import any local module, it returns only the
+    first script mpy data (without size and name) for backwards compatibility
+    with older firmware.
+
+    Arguments:
+        path:
+            Path to script that is to be compiled.
+        abi:
+            Expected MPY ABI version.
+        compile_args:
+            Extra arguments for ``mpy-cross``.
+
+    Returns:
+        Concatenation of all compiled files in the format given above.
+
+    Raises:
+        RuntimeError: if there is not a running event loop.
+        ValueError if MPY ABI version is not 5 or 6.
+        subprocess.CalledProcessError: if executing the ``mpy-cross` tool failed.
+    """
+
+    # Make the build directory
+    make_build_dir()
+
+    # Directory where main and dependencies are located
+    source_dir = Path(path).parent
+
+    # Set of all dependencies
+    dependencies = set()
+    not_found = set()
+
+    # Find all dependencies recursively
+    def find_dependencies(module_name):
+        try:
+            with open(source_dir / (module_name + ".py")) as source:
+                # Search non-recursively through current module
+                local_dependencies = set()
+                for line in source:
+                    # from my_module import thing1, thing2 ---> my_module
+                    if result := re.search("from (.*) import (.*)", line):
+                        local_dependencies.add(result.group(1))
+                    # import my_module ---> my_module
+                    elif result := re.search("import (.*)", line):
+                        local_dependencies.add(result.group(1))
+
+                # If each file that wasn't already done, add it and find its
+                # dependencies too.
+                for dep in local_dependencies.difference(dependencies):
+                    if dep not in dependencies:
+                        dependencies.add(dep)
+                        find_dependencies(dep)
+        # Some modules are stored on the hub so we can't find them here.
+        except FileNotFoundError:
+            not_found.add(module_name)
+
+    # Start searching from the top level.
+    main_module = Path(path).stem
+    find_dependencies(main_module)
+
+    # Subtract the (builtin or missing) modules we won't upload.
+    dependencies = dependencies.difference(not_found)
+
+    print("Uploading:", path)
+    # If there are no dependencies, it is an old-style single file script.
+    # For backwards compatibility, upload just the mpy data. Once the new
+    # firmware is stable, we can remove this special case.
+    if not dependencies:
+        return await compile_file(path, abi)
+
+    # Get the total tuple of main programs and module
+    print("Included modules:", dependencies)
+    modules = [main_module] + sorted(tuple(dependencies))
+
+    # Get a data blob with all scripts.
+    blob = bytearray([])
+    for module in modules:
+        name = module.encode() + b"\x00"
+        mpy = await compile_file(source_dir / (module + ".py"), abi)
+        size = len(mpy).to_bytes(4, "little")
+        blob += size + name + mpy
+    return blob
+
+
 def save_script(py_string):
     """Save a MicroPython one-liner to a file."""
     # Make the build directory.

pybricksdev/connections/pybricks.py

@@ -28,7 +28,7 @@
     StatusFlag,
     unpack_pnp_id,
 )
-from ..compile import compile_file
+from ..compile import compile_multi_file
 from ..tools import chunk
 from ..tools.checksum import xor_bytes
 
@@ -264,7 +264,7 @@ async def run(self, py_path, wait=True, print_output=True):
 
         # Compile the script to mpy format
         self.script_dir, _ = os.path.split(py_path)
-        mpy = await compile_file(py_path, self._mpy_abi_version)
+        mpy = await compile_multi_file(py_path, self._mpy_abi_version)
 
         try:
             self.loading = True