connections: Add USB file upload.

This is used for installing Pybricks on SPIKE hubs without DFU.

If no hub in serial mode is detected, it will still proceed with DFU as it did before.

Author: laurensvalk

pybricksdev/cli/__init__.py

@@ -12,6 +12,7 @@
 from tempfile import NamedTemporaryFile
 from typing import ContextManager, TextIO
 import validators
+import zipfile
 
 from abc import ABC, abstractmethod
 from os import PathLike, path
@@ -229,8 +230,39 @@ async def run(self, args: argparse.Namespace):
 
         if metadata["device-id"] in (HubKind.TECHNIC_SMALL, HubKind.TECHNIC_LARGE):
             from ..dfu import flash_dfu
+            from ..connections import REPLHub
 
-            flash_dfu(firmware, metadata)
+            try:
+                # Connect to the hub and exit the runtime.
+                hub = REPLHub()
+                await hub.connect()
+                await hub.reset_hub()
+
+                # Upload installation script.
+                archive = zipfile.ZipFile(args.firmware)
+                await hub.exec_line("import uos; uos.mkdir('_firmware')")
+                await hub.upload_file(
+                    "_firmware/install_pybricks.py",
+                    bytearray(archive.open("install_pybricks.py").read()),
+                )
+
+                # Upload metadata.
+                await hub.upload_file(
+                    "_firmware/firmware.metadata.json",
+                    bytearray(archive.open("firmware.metadata.json").read()),
+                )
+
+                # Upload Pybricks firmware
+                await hub.upload_file("_firmware/firmware.bin", firmware)
+
+                # Run installation script
+                print("Installing firmware")
+                await hub.exec_line("from _firmware.install_pybricks import install")
+                await hub.exec_paste_mode("install()")
+
+            except OSError:
+                print("Could not find hub in standard firmware mode. Trying DFU.")
+                flash_dfu(firmware, metadata)
         else:
             from ..ble import find_device
             from ..flash import BootloaderConnection

pybricksdev/connections.py

@@ -381,6 +381,47 @@ async def run(self, py_path, wait=True, print_output=True):
             await asyncio.sleep(0.3)
 
 
+FILE_PACKET_SIZE = 1024
+FILE_TRANSFER_SCRIPT = f"""
+import sys
+import micropython
+import utime
+
+PACKETSIZE = {FILE_PACKET_SIZE}
+
+def receive_file(filename, filesize):
+
+    micropython.kbd_intr(-1)
+
+    with open(filename, "wb") as f:
+
+        # Initialize buffers
+        done = 0
+        buf = bytearray(PACKETSIZE)
+        sys.stdin.buffer.read(1)
+
+        while done < filesize:
+
+            # Size of last package
+            if filesize - done < PACKETSIZE:
+                buf = bytearray(filesize - done)
+
+            # Read one packet from standard in.
+            time_now = utime.ticks_ms()
+            bytes_read = sys.stdin.buffer.readinto(buf)
+
+            # If transmission took a long time, something bad happened.
+            if utime.ticks_ms() - time_now > 5000:
+                print("transfer timed out")
+                return
+
+            # Write the data and say we're ready for more.
+            f.write(buf)
+            done += bytes_read
+            print("ACK")
+"""
+
+
 class REPLHub:
     """Run scripts on generic MicroPython boards with a REPL over USB."""
 
@@ -455,6 +496,12 @@ async def reset_hub(self):
         # Clear all buffers
         self.reset_buffers()
 
+        # Load file transfer function
+        await self.exec_paste_mode(FILE_TRANSFER_SCRIPT, print_output=False)
+        self.reset_buffers()
+
+        print("Hub is ready.")
+
     async def exec_line(self, line, wait=True):
         """Executes one line on the REPL."""
 
@@ -467,10 +514,6 @@ async def exec_line(self, line, wait=True):
         echo = encoded + b"\r\n"
         self.serial.write(echo)
 
-        # We are done if we don't want to wait for the result.
-        if not wait:
-            return
-
         # Wait until the echo has been read.
         while len(self.buffer) < start_len + len(echo):
             await asyncio.sleep(0.05)
@@ -480,6 +523,10 @@ async def exec_line(self, line, wait=True):
             print(start_len, self.buffer, self.buffer[start_len - 1 :], echo)
             raise ValueError("Failed to execute line: {0}.".format(line))
 
+        # We are done if we don't want to wait for the result.
+        if not wait:
+            return
+
         # Wait for MicroPython to execute the command.
         while not self.is_idle():
             await asyncio.sleep(0.1)
@@ -556,3 +603,38 @@ async def run(self, py_path, wait=True, print_output=True):
         self.script_dir, _ = os.path.split(py_path)
         await self.reset_hub()
         await self.exec_paste_mode(script, wait, print_output)
+
+    async def upload_file(self, destination, contents):
+        """Uploads a file to the hub."""
+
+        # Print upload info.
+        size = len(contents)
+        print(f"Uploading {destination} ({size} bytes)")
+        self.reset_buffers()
+
+        # Prepare hub to receive file
+        await self.exec_line(f"receive_file('{destination}', {size})", wait=False)
+
+        ACK = b"ACK" + self.EOL
+        progress = 0
+
+        # Write file chunk by chunk.
+        for data in chunk(contents, FILE_PACKET_SIZE):
+
+            # Send a chunk and wait for acknowledgement of receipt
+            buffer_now = len(self.buffer)
+            progress += self.serial.write(data)
+            while len(self.buffer) < buffer_now + len(ACK):
+                await asyncio.sleep(0.01)
+                self.parse_input()
+
+            # Raise error if we didn't get acknowledgement
+            if self.buffer[buffer_now : buffer_now + len(ACK)] != ACK:
+                print(self.buffer[buffer_now:])
+                raise ValueError("Did not get expected response from the hub.")
+
+            # Print progress
+            print(f"Progress: {int(progress / size * 100)}%", end="\r")
+
+        # Get REPL back in normal state
+        await self.exec_line("# File transfer complete")