tools: add chunk method

Chunking data is something we do frequently, so it makes sense to have
a helper function to do this.

Author: dlech

CHANGELOG.md

@@ -90,6 +90,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [1.0.0-alpha.4] - 2021-04-12
 ### Added
 - Size check when restoring firmware via USB/DFU.
+- Added `pybricksdev.tools.chunk()` function.
 ### Fixed
 - Wait for some time to allow program output to be received before disconnecting
   in the `run` command.

pybricksdev/ble/__init__.py

@@ -9,6 +9,7 @@
 from bleak.backends.device import BLEDevice
 from bleak.backends.scanner import AdvertisementData
 
+from ..tools import chunk
 from .pybricks import PYBRICKS_SERVICE_UUID
 
 logger = logging.getLogger(__name__)
@@ -131,16 +132,14 @@ async def write(self, data, with_response=False):
                 Write with or without response.
         """
         # Send the chunks one by one
-        for i in range(0, len(data), self.max_data_size):
-            chunk = data[i : i + self.max_data_size]
+        for c in chunk(data, self.max_data_size):
             logger.debug(
                 "TX CHUNK: {0}, {1} response".format(
-                    chunk, "with" if with_response else "without"
+                    c, "with" if with_response else "without"
                 )
             )
-            # Send one chunk
             await self.client.write_gatt_char(
-                self.char_rx_UUID, bytearray(chunk), with_response
+                self.char_rx_UUID, bytearray(c), with_response
             )
 
 

pybricksdev/compile.py

@@ -9,6 +9,8 @@
 
 import mpy_cross
 
+from .tools import chunk
+
 logger = logging.getLogger(__name__)
 
 BUILD_DIR = "build"
@@ -138,9 +140,8 @@ def print_mpy(data):
     print(f"// MPY file. Version: {data[1]}. Size: {len(data)} bytes")
     print("const uint8_t script[] = {")
 
-    for i in range(0, len(data), WIDTH):
-        chunk = data[i : i + WIDTH]
-        hex_repr = [f"0x{i:02X}" for i in chunk]
+    for c in chunk(data, WIDTH):
+        hex_repr = [f"0x{i:02X}" for i in c]
         print(f"    {', '.join(hex_repr)},")
 
     print("};")

pybricksdev/connections.py

@@ -29,6 +29,7 @@
     unpack_pnp_id,
 )
 from .compile import compile_file
+from .tools import chunk
 from .tools.checksum import xor_bytes
 from .usbconnection import USBConnection
 
@@ -295,17 +296,13 @@ async def run(self, py_path, wait=True, print_output=True):
         length = len(mpy).to_bytes(4, byteorder="little")
         await self.send_message(length)
 
-        # Divide script in chunks of bytes
-        n = 100
-        chunks = [mpy[i : i + n] for i in range(0, len(mpy), n)]
-
         # Send the data chunk by chunk
         with logging_redirect_tqdm(), tqdm(
             total=len(mpy), unit="B", unit_scale=True
         ) as pbar:
-            for chunk in chunks:
-                await self.send_message(chunk)
-                pbar.update(len(chunk))
+            for c in chunk(mpy, 100):
+                await self.send_message(c)
+                pbar.update(len(c))
 
         # Optionally wait for the program to finish
         if wait:
@@ -424,11 +421,6 @@ async def run(self, py_path, wait=False):
         with open(py_path, "rb") as demo:
             program = demo.read()
 
-        chunk_size = 512
-        chunks = [
-            program[i : i + chunk_size] for i in range(0, len(program), chunk_size)
-        ]
-
         while response is None or "transferid" not in response:
             response = await self.send_command_and_get_response(
                 "start_write_program",
@@ -450,15 +442,15 @@ async def run(self, py_path, wait=False):
         with logging_redirect_tqdm(), tqdm(
             total=len(program), unit="B", unit_scale=True
         ) as pbar:
-            for chunk in chunks:
+            for c in chunk(program, 512):
                 response = await self.send_command_and_get_response(
                     "write_package",
                     {
-                        "data": base64.b64encode(chunk).decode("ascii"),
+                        "data": base64.b64encode(c).decode("ascii"),
                         "transferid": transferid,
                     },
                 )
-                pbar.update(len(chunk))
+                pbar.update(len(c))
 
         await asyncio.sleep(0.5)
         response = await self.send_command_and_get_response(
@@ -798,17 +790,13 @@ async def run(self, py_path, wait=True, print_output=True):
             length = len(mpy).to_bytes(4, byteorder="little")
             await self.send_block(length)
 
-            # Divide script in chunks of bytes
-            n = 100
-            chunks = [mpy[i : i + n] for i in range(0, len(mpy), n)]
-
             # Send the data chunk by chunk
             with logging_redirect_tqdm(), tqdm(
                 total=len(mpy), unit="B", unit_scale=True
             ) as pbar:
-                for chunk in chunks:
-                    await self.send_block(chunk)
-                    pbar.update(len(chunk))
+                for c in chunk(mpy, 100):
+                    await self.send_block(c)
+                    pbar.update(len(c))
         finally:
             self.loading = False
 

pybricksdev/tools/__init__.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2021 The Pybricks Authors
+
+from typing import Generator, Sequence, TypeVar
+
+T = TypeVar("T")
+
+
+def chunk(data: Sequence[T], size: int) -> Generator[Sequence[T], None, None]:
+    """
+    Divides a sequence into equal-sized chunks.
+
+    The last chunk may be smaller than *size*.
+
+    Args:
+        data: The data to be divided.
+        size: The size of each chunk.
+
+    Returns:
+        A generator that yields each chunk.
+    """
+    for i in range(0, len(data), size):
+        yield data[i : i + size]
+
+
+__all__ = ["chunk"]

tests/test_chunk.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2021 The Pybricks Authors
+
+
+from pybricksdev.tools import chunk
+
+
+def test_chunk():
+    expected = [
+        [1, 2, 3],
+        [4, 5, 6],
+        [7, 8, 9],
+        [10],
+    ]
+
+    count = 0
+
+    for c, e in zip(chunk([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 3), expected):
+        assert c == e
+        count += 1
+
+    assert count == len(expected)