tools/mpy-tool.py: Allow dumping MPY segments into their own files.

This commit lets "tools/mpy-tool.py" extract MPY segments into their own
files, one file per segment.

A pair of new command line arguments were added, namely "-e"/"--extract"
that takes a filename prefix to use as a base for the generated files'
name, and "--extract-only" that - combined with "--extract" - allows
selecting which kinds of segment should be dumped to the filesystem.

So, for example, assuming there's a file called "module.mpy", running
"./mpy-tool.py --extract segments module.mpy" would yield a series of
files with names like "segments_0_module.py_QSTR_module.py.bin",
"segments_1_module.py_META__module_.bin",
"segments_2_module.py_QSTR_function.bin", etc.  In short the file name
format is "<base>_<count>_<sourcefile>_<segmentkind>_<segmentname>.bin",
with <segmentkind> being META, QSTR, OBJ, or CODE.  Source file names
and segment names will only contain characters in the range
"a-zA-Z0-9_-." to avoid having output file names with unexpected
characters.

The "--extract-only" option can accept one or more kinds, separated by
commas and treated as case insensitive strings.  The supported kinds
match what is currently handled by the "MPYSegment" class in
"tools/mpy-tool.py": "META", "QSTR", "OBJ", and "CODE".  The absence of
this command line option implies dumping every segment found.

If "--extract" is passed along with "--merge", dumping is performed
after the merge process takes place, in order to dump all possible
segments that match the requested segment kinds.

Signed-off-by: Alessandro Gatti <a.gatti@frob.it>

1 files changed, 49 insertions, 0 deletions

diff --git a/tools/mpy-tool.py b/tools/mpy-tool.py
index 69d3a6c06..5c63f5be6 100755
--- a/tools/mpy-tool.py
+++ b/tools/mpy-tool.py
@@ -1765,6 +1765,44 @@ def merge_mpy(compiled_modules, output_file):
             f.write(merged_mpy)
 
 
+def extract_segments(compiled_modules, basename, kinds_arg):
+    import re
+
+    kind_str = ("META", "QSTR", "OBJ", "CODE")
+    kinds = set()
+    if kinds_arg is not None:
+        for kind in kinds_arg.upper().split(","):
+            if kind in kind_str:
+                kinds.add(kind)
+            else:
+                raise Exception('unknown segment kind "%s"' % (kind,))
+    segments = []
+    for module in compiled_modules:
+        for segment in module.mpy_segments:
+            if not kinds or kind_str[segment.kind] in kinds:
+                segments.append((module.mpy_source_file, module.source_file.str, segment))
+    count_len = len(str(len(segments)))
+    sanitiser = re.compile("[^a-zA-Z0-9_.-]")
+    for counter, entry in enumerate(segments):
+        file_name, source_file, segment = entry
+        output_name = (
+            basename
+            + "_"
+            + str(counter).rjust(count_len, "0")
+            + "_"
+            + sanitiser.sub("_", source_file)
+            + "_"
+            + kind_str[segment.kind]
+            + "_"
+            + sanitiser.sub("_", str(segment.name))
+            + ".bin"
+        )
+        with open(file_name, "rb") as source:
+            with open(output_name, "wb") as output:
+                source.seek(segment.start)
+                output.write(source.read(segment.end - segment.start))
+
+
 def main(args=None):
     global global_qstrs
 
@@ -1781,6 +1819,14 @@ def main(args=None):
     cmd_parser.add_argument(
         "--merge", action="store_true", help="merge multiple .mpy files into one"
     )
+    cmd_parser.add_argument(
+        "-e", "--extract", metavar="BASE", type=str, help="write segments into separate files"
+    )
+    cmd_parser.add_argument(
+        "--extract-only",
+        metavar="KIND[,...]",
+        help="extract only segments of the given type (meta, qstr, obj, code)",
+    )
     cmd_parser.add_argument("-q", "--qstr-header", help="qstr header file to freeze against")
     cmd_parser.add_argument(
         "-mlongint-impl",
@@ -1848,6 +1894,9 @@ def main(args=None):
     if args.merge:
         merge_mpy(compiled_modules, args.output)
 
+    if args.extract:
+        extract_segments(compiled_modules, args.extract, args.extract_only)
+
 
 if __name__ == "__main__":
     main()