esp-idf/tools/mkuf2.py

#!/usr/bin/env python
#
# SPDX-FileCopyrightText: 2020-2025 Espressif Systems (Shanghai) CO LTD
# SPDX-License-Identifier: Apache-2.0
# Module was moved to the esptool in ESP-IDF v5.2 and relicensed under GPL v2.0 license.
import argparse
import json
import os
import subprocess
import sys


def main() -> None:
    parser = argparse.ArgumentParser()

    def parse_chip_id(string: str) -> str:
        # compatibility layer with old script
        print("DEPRECATED option '--chip-id'. Please consider using '--chip' instead")
        # DO NOT add new IDs; they are now maintained in esptool.
        ids = {
            0x1C5F21B0: 'esp32',
            0xBFDD4EEE: 'esp32s2',
            0xD42BA06C: 'esp32c3',
            0xC47E5767: 'esp32s3',
            0x332726F6: 'esp32h2',
            0x2B88D29C: 'esp32c2',
            0x540DDF62: 'esp32c6',
            0x3D308E94: 'esp32p4',
        }
        try:
            return ids[int(string, 16)]
        except KeyError:
            raise argparse.ArgumentTypeError('Unknown Chip ID')

    # Provision to add "info" command
    subparsers = parser.add_subparsers(dest='command')
    write_parser = subparsers.add_parser('write')
    write_parser.add_argument('-o', '--output-file', help='Filename for storing the output UF2 image', required=True)
    group = write_parser.add_mutually_exclusive_group(required=True)
    # chip-id used just for backwards compatibility, UF2 family IDs are now stored in esptool
    group.add_argument('--chip-id', type=parse_chip_id, help=argparse.SUPPRESS)
    group.add_argument('--chip', type=str, help='Target chip type')
    write_parser.add_argument(
        '--chunk-size',
        required=False,
        type=int,
        default=None,
        help='Specify the used data part of the 512 byte UF2 block. A common value is 256. By '
        'default the largest possible value will be used.',
    )
    write_parser.add_argument(
        '--json', help='Optional file for loading "flash_files" dictionary with <address> <file> items'
    )
    write_parser.add_argument(
        '--bin',
        help='Use only a subset of binaries from the JSON file, e.g. "partition_table bootloader app"',
        nargs='*',
    )
    write_parser.add_argument(
        '--md5-disable', help='Disable MD5 checksum. Useful for compatibility with e.g. TinyUF2', action='store_true'
    )
    write_parser.add_argument('files', metavar='<address> <file>', help='Add <file> at <address>', nargs='*')

    args = parser.parse_args()

    def check_file(file_name: str) -> str:
        if not os.path.isfile(file_name):
            raise RuntimeError(f'{file_name} is not a regular file!')
        return file_name

    files: list[tuple[int, str]] = []
    if args.files:
        files += [(addr, check_file(f_name)) for addr, f_name in zip(args.files[::2], args.files[1::2])]

    if args.json:
        json_dir = os.path.dirname(os.path.abspath(args.json))

        def process_json_file(path: str) -> str:
            """
            The input path is relative to json_dir. This function makes it relative to the current working
            directory.
            """
            return check_file(os.path.abspath(os.path.join(json_dir, path)))

        with open(args.json, encoding='utf-8') as f:
            json_content = json.load(f)

        if args.bin:
            try:
                bin_selection = [json_content[b] for b in args.bin]
                flash_dic = dict((x['offset'], x['file']) for x in bin_selection)
            except KeyError:
                print('Invalid binary was selected.')
                valid = [k if all(x in v for x in ('offset', 'file')) else None for k, v in json_content.items()]
                print('Valid ones:', ' '.join(x for x in valid if x))
                exit(1)
        else:
            flash_dic = json_content['flash_files']

        files += [(addr, process_json_file(f_name)) for addr, f_name in flash_dic.items()]

    # remove possible duplicates and sort based on the address
    files = sorted([(addr, f_name) for addr, f_name in dict(files).items()], key=lambda x: x[0])

    # list of tuples to simple list
    files_flatten = [item for t in files for item in t]

    cmd = [
        sys.executable,
        '-m',
        'esptool',
        '--chip',
        args.chip_id or args.chip,
        'merge-bin',
        '--format',
        'uf2',
        '-o',
        args.output_file,
    ]
    if args.chunk_size:
        cmd.extend(['--chunk-size', args.chunk_size])

    if args.md5_disable:
        cmd.append('--md5-disable')

    cmd_str = ' '.join(cmd + files_flatten)
    print(f'Executing: {cmd_str}')

    sys.exit(subprocess.run(cmd + files_flatten).returncode)


if __name__ == '__main__':
    main()