parttool, otatool: accept esptool args

components/partition_table/parttool.py

@@ -22,6 +22,7 @@ import os
 import sys
 import subprocess
 import tempfile
+import re
 import gen_esp32part as gen
 
 
@@ -66,11 +67,30 @@ PARTITION_BOOT_DEFAULT = _PartitionId()
 
 class ParttoolTarget():
 
-    def __init__(self, port=None, partition_table_offset=PARTITION_TABLE_OFFSET, partition_table_file=None):
+    def __init__(self, port=None, partition_table_offset=PARTITION_TABLE_OFFSET, partition_table_file=None,
+                 esptool_args=[], esptool_write_args=[], esptool_read_args=[], esptool_erase_args=[]):
         self.port = port
 
         gen.offset_part_table = partition_table_offset
 
+        def parse_esptool_args(esptool_args):
+            results = list()
+            for arg in esptool_args:
+                pattern = re.compile(r"(.+)=(.+)")
+                result = pattern.match(arg)
+                try:
+                    key = result.group(1)
+                    value = result.group(2)
+                    results.extend(["--" + key, value])
+                except AttributeError:
+                    results.extend(["--" + arg])
+            return results
+
+        self.esptool_args = parse_esptool_args(esptool_args)
+        self.esptool_write_args = parse_esptool_args(esptool_write_args)
+        self.esptool_read_args = parse_esptool_args(esptool_read_args)
+        self.esptool_erase_args = parse_esptool_args(esptool_erase_args)
+
         if partition_table_file:
             try:
                 with open(partition_table_file, "rb") as f:
@@ -93,7 +113,7 @@ class ParttoolTarget():
         self.partition_table = partition_table
 
     def _call_esptool(self, args, out=None):
-        esptool_args = [sys.executable, ESPTOOL_PY]
+        esptool_args = [sys.executable, ESPTOOL_PY] + self.esptool_args
 
         if self.port:
             esptool_args += ["--port", self.port]
@@ -124,11 +144,11 @@ class ParttoolTarget():
 
     def erase_partition(self, partition_id):
         partition = self.get_partition_info(partition_id)
-        self._call_esptool(["erase_region", str(partition.offset),  str(partition.size)])
+        self._call_esptool(["erase_region", str(partition.offset),  str(partition.size)] + self.esptool_erase_args)
 
     def read_partition(self, partition_id, output):
         partition = self.get_partition_info(partition_id)
-        self._call_esptool(["read_flash", str(partition.offset), str(partition.size), output])
+        self._call_esptool(["read_flash", str(partition.offset), str(partition.size), output] + self.esptool_read_args)
 
     def write_partition(self, partition_id, input):
         self.erase_partition(partition_id)
@@ -141,7 +161,7 @@ class ParttoolTarget():
             if content_len > partition.size:
                 raise Exception("Input file size exceeds partition size")
 
-        self._call_esptool(["write_flash", str(partition.offset), input])
+        self._call_esptool(["write_flash", str(partition.offset), input] + self.esptool_write_args)
 
 
 def _write_partition(target, partition_id, input):
@@ -191,6 +211,10 @@ def main():
     parser = argparse.ArgumentParser("ESP-IDF Partitions Tool")
 
     parser.add_argument("--quiet", "-q", help="suppress stderr messages", action="store_true")
+    parser.add_argument("--esptool-args", help="additional main arguments for esptool", nargs="+")
+    parser.add_argument("--esptool-write-args", help="additional subcommand arguments when writing to flash", nargs="+")
+    parser.add_argument("--esptool-read-args", help="additional subcommand arguments when reading flash", nargs="+")
+    parser.add_argument("--esptool-erase-args", help="additional subcommand arguments when erasing regions of flash", nargs="+")
 
     # By default the device attached to the specified port is queried for the partition table. If a partition table file
     # is specified, that is used instead.
@@ -264,6 +288,18 @@ def main():
     if args.partition_table_offset:
         target_args["partition_table_offset"] = int(args.partition_table_offset, 0)
 
+    if args.esptool_args:
+        target_args["esptool_args"] = args.esptool_args
+
+    if args.esptool_write_args:
+        target_args["esptool_write_args"] = args.esptool_write_args
+
+    if args.esptool_read_args:
+        target_args["esptool_read_args"] = args.esptool_read_args
+
+    if args.esptool_erase_args:
+        target_args["esptool_erase_args"] = args.esptool_erase_args
+
     target = ParttoolTarget(**target_args)
 
     # Create the operation table and execute the operation