diff --git a/components/partition_table/check_sizes.py b/components/partition_table/check_sizes.py index 2045155143e..489de810070 100755 --- a/components/partition_table/check_sizes.py +++ b/components/partition_table/check_sizes.py @@ -7,15 +7,12 @@ # # (Can also check if the bootloader binary fits before the partition table.) # -# SPDX-FileCopyrightText: 2020-2022 Espressif Systems (Shanghai) CO LTD +# SPDX-FileCopyrightText: 2020-2026 Espressif Systems (Shanghai) CO LTD # SPDX-License-Identifier: Apache-2.0 -from __future__ import division, print_function, unicode_literals -import argparse import io # noqa: F401 # pylint: disable=unused-import import os -import sys try: from typing import IO # noqa: F401 # pylint: disable=unused-import @@ -23,12 +20,17 @@ except ImportError: pass # used for type hinting only import gen_esp32part -from gen_esp32part import PartitionTable, get_ptype_as_int, get_subtype_as_int +import rich_click as click +from esp_pylib.cli_types import AnyIntType +from esp_pylib.logger import log +from gen_esp32part import PartitionTable +from gen_esp32part import get_ptype_as_int +from gen_esp32part import get_subtype_as_int allow_failures = False -def _file_size(f): # type: (IO) -> int +def _file_size(f): # type: (IO) -> int before = f.tell() f.seek(0, 2) # seek to end result = f.tell() @@ -36,25 +38,27 @@ def _file_size(f): # type: (IO) -> int return result -def _fail(msg): # type: (str) -> None +def _fail(msg): # type: (str) -> None if allow_failures: - print('Warning: {}'.format(msg)) + log.warn(msg) else: - raise SystemExit('Error: {}'.format(msg)) + log.err(msg) + raise SystemExit(msg) def check_bootloader(partition_table_offset, bootloader_offset, binary_file): # type: (int, int, IO) -> None max_size = partition_table_offset - bootloader_offset bootloader_size = _file_size(binary_file) if bootloader_size > max_size: - msg = ('Bootloader binary size {:#x} bytes is too large for partition table offset {:#02x}. ' + - 'Bootloader binary can be maximum {:#x} ({}) bytes unless the partition table offset ' + - 'is increased in the Partition Table section of the project configuration menu.').format( - bootloader_size, partition_table_offset, max_size, max_size) + msg = ( + 'Bootloader binary size {:#x} bytes is too large for partition table offset {:#02x}. ' + + 'Bootloader binary can be maximum {:#x} ({}) bytes unless the partition table offset ' + + 'is increased in the Partition Table section of the project configuration menu.' + ).format(bootloader_size, partition_table_offset, max_size, max_size) _fail(msg) free_size = max_size - bootloader_size - print('Bootloader binary size {:#x} bytes. {:#x} bytes ({}%) free.'.format( - bootloader_size, free_size, round(free_size * 100 / max_size))) + pct = round(free_size * 100 / max_size) + log.print(f'Bootloader binary size {bootloader_size:#x} bytes. {free_size:#x} bytes ({pct}%) free.') def check_partition(ptype, subtype, partition_table_file, bin_file): # type: (str, str, io.IOBase, IO) -> None @@ -64,12 +68,12 @@ def check_partition(ptype, subtype, partition_table_file, bin_file): # type: (s partitions = [p for p in table if p.type == ptype] if subtype is not None: - ptype_str += ' ({})'.format(subtype) + ptype_str += f' ({subtype})' subtype = get_subtype_as_int(ptype, subtype) partitions = [p for p in partitions if p.subtype == subtype] if len(partitions) == 0: - print('WARNING: Partition table does not contain any partitions matching {}'.format(ptype_str)) + log.warn(f'Partition table does not contain any partitions matching {ptype_str}') return bin_name = os.path.basename(bin_file.name) @@ -78,65 +82,87 @@ def check_partition(ptype, subtype, partition_table_file, bin_file): # type: (s if smallest_size >= bin_size: free_size = smallest_size - bin_size free_size_relative = free_size / smallest_size - print('{} binary size {:#x} bytes. Smallest {} partition is {:#x} bytes. {:#x} bytes ({:.0%}) free.'.format( - bin_name, bin_size, ptype_str, smallest_size, free_size, free_size_relative)) + log.print( + f'{bin_name} binary size {bin_size:#x} bytes. Smallest {ptype_str} partition is ' + f'{smallest_size:#x} bytes. {free_size:#x} bytes ({free_size_relative:.0%}) free.' + ) free_size_relative_critical = 0.05 if free_size_relative < free_size_relative_critical: - print('Warning: The smallest {} partition is nearly full ({:.0%} free space left)!'.format(ptype_str, free_size_relative)) + log.warn(f'The smallest {ptype_str} partition is nearly full ({free_size_relative:.0%} free space left)!') return too_small_partitions = [p for p in partitions if p.size < bin_size] if len(partitions) == 1: - msg = '{} partition is'.format(ptype_str) + msg = f'{ptype_str} partition is' elif len(partitions) == len(too_small_partitions): - msg = 'All {} partitions are'.format(ptype_str) + msg = f'All {ptype_str} partitions are' else: - msg = '{}/{} {} partitions are'.format(len(too_small_partitions), len(partitions), ptype_str) - msg += ' too small for binary {} size {:#x}:'.format(bin_name, bin_size) + msg = f'{len(too_small_partitions)}/{len(partitions)} {ptype_str} partitions are' + msg += f' too small for binary {bin_name} size {bin_size:#x}:' for p in too_small_partitions: - msg += '\n - {} (overflow {:#x})'.format(p, bin_size - p.size) + msg += f'\n - {p} (overflow {bin_size - p.size:#x})' if not allow_failures and len(partitions) == len(too_small_partitions): # if some partitions can fit the binary then just print a warning - raise SystemExit('Error: ' + msg) + _fail(msg) else: - print('Warning: ' + msg) + log.warn(msg) + + +def _configure_globals(offset: int, allow_failures_flag: bool) -> None: + global allow_failures # pylint: disable=global-statement + + gen_esp32part.quiet = True + gen_esp32part.offset_part_table = offset + allow_failures = allow_failures_flag + + +@click.group( + context_settings={'help_option_names': ['-h', '--help']}, + help='Check binary sizes against partition table entries', +) +@click.option('--target', type=click.Choice(['esp32', 'esp32s2'])) +@click.option( + '--allow-failures', + 'allow_failures_flag', + is_flag=True, + help='If true, failures will print warnings but not exit with an error', +) +@click.option('--offset', '-o', type=AnyIntType(), default=0x8000, show_default=True, help='Set partition table offset') +@click.pass_context +def cli(ctx, target, allow_failures_flag, offset): # type: ignore[no-untyped-def] + ctx.ensure_object(dict) + ctx.obj['offset'] = offset + ctx.obj['allow_failures'] = allow_failures_flag + + +@cli.command('bootloader', help='Check bootloader binary size against partition table offset') +@click.argument('bootloader_offset', type=AnyIntType()) +@click.argument('bootloader_binary', type=click.File('rb')) +@click.pass_context +def bootloader_cmd(ctx, bootloader_offset, bootloader_binary): # type: ignore[no-untyped-def] + _configure_globals(ctx.obj['offset'], ctx.obj['allow_failures']) + check_bootloader(ctx.obj['offset'], bootloader_offset, bootloader_binary) + + +@cli.command('partition', help='Check binary size against partition table entries') +@click.option( + '--type', 'partition_type', required=True, help='Check the file size against all partitions of this type.' +) +@click.option('--subtype', help='Optional, only check the file size against all partitions of this subtype.') +@click.argument('partition_table', type=click.File('rb')) +@click.argument('binary', type=click.File('rb')) +@click.pass_context +def partition_cmd(ctx, partition_type, subtype, partition_table, binary): # type: ignore[no-untyped-def] + _configure_globals(ctx.obj['offset'], ctx.obj['allow_failures']) + check_partition(partition_type, subtype, partition_table, binary) def main(): # type: () -> None - global allow_failures # pylint: disable=global-statement - - parser = argparse.ArgumentParser(description='Check binary sizes against partition table entries') - parser.add_argument('--target', choices=['esp32', 'esp32s2']) - parser.add_argument('--allow_failures', action='store_true', help='If true, failures will print warnings but not exit with an error') - parser.add_argument('--offset', '-o', help='Set partition table offset', default='0x8000') - - subparsers = parser.add_subparsers(dest='check_target', - help='Type of binary to check against partition table layout') - sp_bootloader = subparsers.add_parser('bootloader') - sp_bootloader.add_argument('bootloader_offset', help='Hex offset of bootloader in flash') - sp_bootloader.add_argument('bootloader_binary', type=argparse.FileType('rb'), help='Bootloader binary (.bin) file from build output') - - sp_part = subparsers.add_parser('partition') - sp_part.add_argument('--type', type=str, help='Check the file size against all partitions of this type.', required=True) - sp_part.add_argument('--subtype', type=str, help='Optional, only check the file size against all partitions of this subtype.') - sp_part.add_argument('partition_table', type=argparse.FileType('rb'), help='Partition table file') - sp_part.add_argument('binary', type=argparse.FileType('rb'), help='Binary file which will have the size checked') - - args = parser.parse_args() - - gen_esp32part.quiet = True - - args.offset = int(args.offset, 0) - gen_esp32part.offset_part_table = args.offset - - if args.check_target is None: # add_subparsers only has a 'required' argument since Python 3 - parser.print_help() - sys.exit(1) - if args.check_target == 'bootloader': - check_bootloader(args.offset, int(args.bootloader_offset, 0), args.bootloader_binary) - else: - check_partition(args.type, args.subtype, args.partition_table, args.binary) + cli() if __name__ == '__main__': + from esp_pylib.excepthook import install_exception_reporting + + install_exception_reporting() main() diff --git a/components/partition_table/gen_empty_partition.py b/components/partition_table/gen_empty_partition.py index a49217f78da..240ec6aa24d 100644 --- a/components/partition_table/gen_empty_partition.py +++ b/components/partition_table/gen_empty_partition.py @@ -4,21 +4,20 @@ # # This tool generates an empty binary file of the required size. # -# SPDX-FileCopyrightText: 2018-2022 Espressif Systems (Shanghai) CO LTD +# SPDX-FileCopyrightText: 2018-2026 Espressif Systems (Shanghai) CO LTD # SPDX-License-Identifier: Apache-2.0 -from __future__ import division, print_function, unicode_literals -import argparse import sys +import rich_click as click +from esp_pylib.cli_types import AnyIntType + __version__ = '1.0' -quiet = False - def generate_blanked_file(size, output_path): - output = b'\xFF' * size + output = b'\xff' * size try: stdout_binary = sys.stdout.buffer # Python 3 except AttributeError: @@ -27,28 +26,23 @@ def generate_blanked_file(size, output_path): f.write(output) -def main(): - parser = argparse.ArgumentParser(description='Generates an empty binary file of the required size.') - parser.add_argument('size', help='Size of generated the file', type=str) - - parser.add_argument('output', help='Path for binary file.', nargs='?', default='-') - args = parser.parse_args() - - size = int(args.size, 0) +@click.command( + context_settings={'help_option_names': ['-h', '--help']}, + help='Generates an empty binary file of the required size.', +) +@click.argument('size', type=AnyIntType(), metavar='SIZE') +@click.argument('output', type=click.Path(), required=False, default='-') +def cli(size, output): if size > 0: - generate_blanked_file(size, args.output) - return 0 + generate_blanked_file(size, output) -class InputError(RuntimeError): - def __init__(self, e): - super(InputError, self).__init__(e) +def main(): + cli() if __name__ == '__main__': - try: - r = main() - sys.exit(r) - except InputError as e: - print(e, file=sys.stderr) - sys.exit(2) + from esp_pylib.excepthook import install_exception_reporting + + install_exception_reporting() + main() diff --git a/components/partition_table/gen_esp32part.py b/components/partition_table/gen_esp32part.py index 53790846d5c..97e63ccf4b1 100755 --- a/components/partition_table/gen_esp32part.py +++ b/components/partition_table/gen_esp32part.py @@ -7,9 +7,9 @@ # See https://docs.espressif.com/projects/esp-idf/en/latest/api-guides/partition-tables.html # for explanation of partition table structure and uses. # -# SPDX-FileCopyrightText: 2016-2025 Espressif Systems (Shanghai) CO LTD +# SPDX-FileCopyrightText: 2016-2026 Espressif Systems (Shanghai) CO LTD # SPDX-License-Identifier: Apache-2.0 -import argparse +# ruff: noqa: UP031, UP008 import binascii import codecs import errno @@ -19,6 +19,10 @@ import re import struct import sys +import rich_click as click +from esp_pylib.cli_types import AnyIntType +from esp_pylib.logger import log + MAX_PARTITION_LENGTH = 0xC00 # 3K for partition data (96 entries) leaves 1K in a 4K sector for signature MD5_PARTITION_BEGIN = b'\xeb\xeb' + b'\xff' * 14 # The first 2 bytes are like magic numbers for MD5 sum PARTITION_TABLE_SIZE = 0x1000 # Size of partition table @@ -167,7 +171,8 @@ recovery_bootloader_offset = None def status(msg): """Print status message to stderr""" if not quiet: - critical(msg) + sys.stderr.write(msg) + sys.stderr.write('\n') def critical(msg): @@ -420,7 +425,7 @@ class PartitionTable(list): return '\n'.join(rows) + '\n' -class PartitionDefinition(object): +class PartitionDefinition: MAGIC_BYTES = b'\xaa\x50' # dictionary maps flag name (as used in CSV flags list, property name) @@ -689,82 +694,54 @@ def parse_int(v, keywords={}): raise InputError("Value '%s' is not valid. Known keywords: %s" % (v, ', '.join(keywords))) -def main(): +def run_gen( + input_file, + output='-', + flash_size=None, + disable_md5sum=False, + no_verify=False, + quiet_flag=False, + offset=0x8000, + primary_bootloader_offset_arg=None, + recovery_bootloader_offset_arg=None, + secure_arg=None, + extra_partition_subtypes=None, +): global quiet global md5sum global offset_part_table global secure global primary_bootloader_offset global recovery_bootloader_offset - parser = argparse.ArgumentParser(description='ESP32 partition table utility') - parser.add_argument( - '--flash-size', - help='Optional flash size limit, checks partition table fits in flash', - nargs='?', - choices=['1MB', '2MB', '4MB', '8MB', '16MB', '32MB', '64MB', '128MB'], - ) - parser.add_argument( - '--disable-md5sum', help='Disable md5 checksum for the partition table', default=False, action='store_true' - ) - parser.add_argument('--no-verify', help="Don't verify partition table fields", action='store_true') - parser.add_argument( - '--verify', - '-v', - help='Verify partition table fields (deprecated, this behaviour is ' - 'enabled by default and this flag does nothing.', - action='store_true', - ) - parser.add_argument('--quiet', '-q', help="Don't print non-critical status messages to stderr", action='store_true') - parser.add_argument('--offset', '-o', help='Set offset partition table', default='0x8000') - parser.add_argument('--primary-bootloader-offset', help='Set primary bootloader offset', default=None) - parser.add_argument('--recovery-bootloader-offset', help='Set recovery bootloader offset', default=None) - parser.add_argument( - '--secure', - help='Require app partitions to be suitable for secure boot', - nargs='?', - const=SECURE_V1, - choices=[SECURE_V1, SECURE_V2], - ) - parser.add_argument('--extra-partition-subtypes', help='Extra partition subtype entries', nargs='*') - parser.add_argument('input', help='Path to CSV or binary file to parse.', type=argparse.FileType('rb')) - parser.add_argument( - 'output', - help='Path to output converted binary or CSV file. Will use stdout if omitted.', - nargs='?', - default='-', - ) - - args = parser.parse_args() - - quiet = args.quiet - md5sum = not args.disable_md5sum - secure = args.secure - offset_part_table = int(args.offset, 0) - if args.primary_bootloader_offset is not None: - primary_bootloader_offset = int(args.primary_bootloader_offset, 0) + quiet = quiet_flag + md5sum = not disable_md5sum + secure = secure_arg + offset_part_table = offset + if primary_bootloader_offset_arg is not None: + primary_bootloader_offset = primary_bootloader_offset_arg if primary_bootloader_offset >= offset_part_table: raise InputError( f'Unsupported configuration. Primary bootloader must be below partition table. ' f'Check --primary-bootloader-offset={primary_bootloader_offset:#x} and --offset={offset_part_table:#x}' ) - if args.recovery_bootloader_offset is not None: - recovery_bootloader_offset = int(args.recovery_bootloader_offset, 0) - if args.extra_partition_subtypes: - add_extra_subtypes(args.extra_partition_subtypes) + if recovery_bootloader_offset_arg is not None: + recovery_bootloader_offset = recovery_bootloader_offset_arg + if extra_partition_subtypes: + add_extra_subtypes(extra_partition_subtypes) - table, input_is_binary = PartitionTable.from_file(args.input) + table, input_is_binary = PartitionTable.from_file(input_file) - if not args.no_verify: + if not no_verify: status('Verifying table...') table.verify() - if args.flash_size: - size_mb = int(args.flash_size.replace('MB', '')) + if flash_size: + size_mb = int(flash_size.replace('MB', '')) table.verify_size_fits(size_mb * 1024 * 1024) # Make sure that the output directory is created - output_dir = os.path.abspath(os.path.dirname(args.output)) + output_dir = os.path.abspath(os.path.dirname(output)) if not os.path.exists(output_dir): try: @@ -774,17 +751,82 @@ def main(): raise if input_is_binary: - output = table.to_csv() - with sys.stdout if args.output == '-' else open(args.output, 'w', encoding='utf-8') as f: - f.write(output) + output_data = table.to_csv() + with sys.stdout if output == '-' else open(output, 'w', encoding='utf-8') as f: + f.write(output_data) else: - output = table.to_binary() + output_data = table.to_binary() try: stdout_binary = sys.stdout.buffer # Python 3 except AttributeError: stdout_binary = sys.stdout - with stdout_binary if args.output == '-' else open(args.output, 'wb') as f: - f.write(output) + with stdout_binary if output == '-' else open(output, 'wb') as f: + f.write(output_data) + + +@click.command( + context_settings={'help_option_names': ['-h', '--help']}, + help='ESP32 partition table utility', +) +@click.option( + '--flash-size', + type=click.Choice(['1MB', '2MB', '4MB', '8MB', '16MB', '32MB', '64MB', '128MB']), + help='Optional flash size limit, checks partition table fits in flash', +) +@click.option('--disable-md5sum', is_flag=True, help='Disable md5 checksum for the partition table') +@click.option('--no-verify', is_flag=True, help="Don't verify partition table fields") +@click.option( + '--verify', + '-v', + is_flag=True, + help='Verify partition table fields (deprecated, this behaviour is enabled by default and this flag does nothing).', +) +@click.option('--quiet', '-q', is_flag=True, help="Don't print non-critical status messages to stderr") +@click.option('--offset', '-o', type=AnyIntType(), default=0x8000, show_default=True, help='Set offset partition table') +@click.option('--primary-bootloader-offset', type=AnyIntType(), default=None, help='Set primary bootloader offset') +@click.option('--recovery-bootloader-offset', type=AnyIntType(), default=None, help='Set recovery bootloader offset') +@click.option( + '--secure', + is_flag=False, + flag_value=SECURE_V1, + type=click.Choice([SECURE_V1, SECURE_V2], case_sensitive=False), + default=None, + help='Require app partitions to be suitable for secure boot', +) +@click.option('--extra-partition-subtypes', multiple=True, help='Extra partition subtype entries') +@click.argument('input_file', type=click.File('rb'), metavar='INPUT') +@click.argument('output', type=click.Path(), required=False, default='-') +def cli( + input_file, + output, + flash_size, + disable_md5sum, + no_verify, + verify, + quiet, + offset, + primary_bootloader_offset, + recovery_bootloader_offset, + secure, + extra_partition_subtypes, +): + run_gen( + input_file=input_file, + output=output, + flash_size=flash_size, + disable_md5sum=disable_md5sum, + no_verify=no_verify, + quiet_flag=quiet, + offset=offset, + primary_bootloader_offset_arg=primary_bootloader_offset, + recovery_bootloader_offset_arg=recovery_bootloader_offset, + secure_arg=secure, + extra_partition_subtypes=extra_partition_subtypes or None, + ) + + +def main(): + cli() class InputError(RuntimeError): @@ -798,8 +840,12 @@ class ValidationError(InputError): if __name__ == '__main__': + from esp_pylib.excepthook import install_exception_reporting + + install_exception_reporting() try: main() except InputError as e: - print(e, file=sys.stderr) + # soft_wrap avoids Rich line-breaking long errors (build-system tests grep substrings). + log.print(str(e), file=sys.stderr, soft_wrap=True, markup=False) sys.exit(2) diff --git a/components/partition_table/gen_extra_subtypes_inc.py b/components/partition_table/gen_extra_subtypes_inc.py index 59814bfa2fa..6f640c8513d 100755 --- a/components/partition_table/gen_extra_subtypes_inc.py +++ b/components/partition_table/gen_extra_subtypes_inc.py @@ -1,7 +1,8 @@ #!/usr/bin/env python -# SPDX-FileCopyrightText: 2022 Espressif Systems (Shanghai) CO LTD +# SPDX-FileCopyrightText: 2022-2026 Espressif Systems (Shanghai) CO LTD # SPDX-License-Identifier: Apache-2.0 -import argparse +import rich_click as click +from esp_pylib.logger import log def gen_header_file(path: str, subtypes: str) -> None: @@ -22,15 +23,27 @@ def gen_header_file(path: str, subtypes: str) -> None: fields[0] = fields[0].strip() fields[1] = fields[1].strip() fields[2] = fields[2].strip() - f.write('ESP_PARTITION_SUBTYPE_%s_%s = %s,\n' % (fields[0].upper(), fields[1].upper(), fields[2])) + f.write(f'ESP_PARTITION_SUBTYPE_{fields[0].upper()}_{fields[1].upper()} = {fields[2]},\n') except ValueError as err: - raise ValueError('Error parsing custom subtypes: %s' % err) + log.die(f'Error parsing custom subtypes: {err}') + + +@click.command( + context_settings={'help_option_names': ['-h', '--help']}, + help='ESP32 extra partitions utility', +) +@click.argument('config_dir', type=click.Path()) +@click.argument('extra_partition_subtypes', nargs=-1) +def cli(config_dir, extra_partition_subtypes): # type: ignore[no-untyped-def] + gen_header_file(config_dir, extra_partition_subtypes) + + +def main() -> None: + cli() if __name__ == '__main__': - parser = argparse.ArgumentParser(description='ESP32 extra partitions utility') - parser.add_argument('config_dir', help='Path to config directory') - parser.add_argument('extra_partition_subtypes', help='Extra partition subtype entries', nargs='*') - args = parser.parse_args() + from esp_pylib.excepthook import install_exception_reporting - gen_header_file(args.config_dir, args.extra_partition_subtypes) + install_exception_reporting() + main() diff --git a/components/partition_table/parttool.py b/components/partition_table/parttool.py index 8f2bcdf2c4d..adf98a632b1 100755 --- a/components/partition_table/parttool.py +++ b/components/partition_table/parttool.py @@ -3,9 +3,8 @@ # parttool is used to perform partition level operations - reading, # writing, erasing and getting info about the partition. # -# SPDX-FileCopyrightText: 2018-2025 Espressif Systems (Shanghai) CO LTD +# SPDX-FileCopyrightText: 2018-2026 Espressif Systems (Shanghai) CO LTD # SPDX-License-Identifier: Apache-2.0 -import argparse import os import re import subprocess @@ -13,6 +12,12 @@ import sys import tempfile import gen_esp32part as gen +import rich_click as click +from esp_pylib.cli_options import EspRichGroup +from esp_pylib.cli_options import MutuallyExclusiveOption +from esp_pylib.cli_options import OptionEatAll +from esp_pylib.cli_types import AnyIntType +from esp_pylib.logger import log __version__ = '2.2' @@ -21,12 +26,28 @@ COMPONENTS_PATH = os.path.expandvars(os.path.join('$IDF_PATH', 'components')) PARTITION_TABLE_OFFSET = 0x8000 +class _StringListParamType(click.ParamType): + """ParamType for `OptionEatAll` without ``multiple=True`` (argparse ``nargs='+'``).""" + + name = 'text_list' + + def convert(self, value, param, ctx): + if not value: + return () + if isinstance(value, (list, tuple)): # noqa: UP038 + return tuple(value) + return (value,) + + +_ESPTOOL_ARGS_TYPE = _StringListParamType() + + quiet = False def status(msg): if not quiet: - print(msg) + log.print(msg, file=sys.stderr) class _PartitionId: @@ -122,11 +143,11 @@ class ParttoolTarget: esptool_args += args - print(f'Running {" ".join(esptool_args)}...') + log.print(f'Running {" ".join(esptool_args)}...', file=sys.stderr) try: subprocess.check_call(esptool_args, stdout=out, stderr=subprocess.STDOUT) except subprocess.CalledProcessError as e: - print('An exception: **', str(e), '** occurred in _call_esptool.', file=out) + log.print(f'An exception: ** {e} ** occurred in _call_esptool.', file=out or sys.stderr) raise e def get_partition_info(self, partition_id): @@ -162,7 +183,7 @@ class ParttoolTarget: partition = self.get_partition_info(partition_id) if partition.readonly and not ignore_readonly: - raise SystemExit(f'"{partition.name}" partition is read-only, (use the --ignore-readonly flag to skip it)') + log.die(f'"{partition.name}" partition is read-only, (use the --ignore-readonly flag to skip it)') self.erase_partition(partition_id) @@ -219,165 +240,84 @@ def _get_partition_info(target, partition_id, info): for i in info: infos += [info_dict[i]] except KeyError: - raise RuntimeError(f'Request for unknown partition info {i}') + log.die(f'Request for unknown partition info {i}') print(' '.join(infos)) -def main(): +def _partition_selection_options(func): + decorators = [ + click.option( + '--partition-name', + '-n', + cls=MutuallyExclusiveOption, + exclusive_with=['partition_type', 'partition_boot_default'], + help='name of the partition', + ), + click.option( + '--partition-type', + '-t', + cls=MutuallyExclusiveOption, + exclusive_with=['partition_name', 'partition_boot_default'], + help='type of the partition', + ), + click.option( + '--partition-boot-default', + '-d', + is_flag=True, + cls=MutuallyExclusiveOption, + exclusive_with=['partition_name', 'partition_type'], + help='select the default boot partition using the same fallback logic as the IDF bootloader', + ), + click.option('--partition-subtype', '-s', help='subtype of the partition'), + click.option('--extra-partition-subtypes', multiple=True, help='Extra partition subtype entries'), + ] + for decorator in reversed(decorators): + func = decorator(func) + return func + + +def _build_partition_id(partition_name, partition_type, partition_subtype, partition_boot_default, part_list=False): + if partition_name: + return PartitionName(partition_name) + if partition_type: + if not partition_subtype: + log.die('--partition-subtype should be defined when --partition-type is defined') + return PartitionType(partition_type, partition_subtype, part_list) + if partition_boot_default: + return PARTITION_BOOT_DEFAULT + log.die( + 'Partition to operate on should be defined using --partition-name OR ' + 'partition-type,--partition-subtype OR partition-boot-default' + ) + + +def _target_kwargs_from_ctx(ctx_obj, extra_partition_subtypes=None): + kwargs = {} + for key, value in ( + ('port', ctx_obj.get('port')), + ('baud', ctx_obj.get('baud')), + ('partition_table_offset', ctx_obj.get('partition_table_offset')), + ('primary_bootloader_offset', ctx_obj.get('primary_bootloader_offset')), + ('recovery_bootloader_offset', ctx_obj.get('recovery_bootloader_offset')), + ('partition_table_file', ctx_obj.get('partition_table_file')), + ('esptool_args', ctx_obj.get('esptool_args')), + ('esptool_write_args', ctx_obj.get('esptool_write_args')), + ('esptool_read_args', ctx_obj.get('esptool_read_args')), + ('esptool_erase_args', ctx_obj.get('esptool_erase_args')), + ): + if value is not None and value != (): + kwargs[key] = value + if extra_partition_subtypes: + gen.add_extra_subtypes(extra_partition_subtypes) + return kwargs + + +def _run_operation(operation, quiet_flag, target, partition_id, **op_kwargs): global quiet - parser = argparse.ArgumentParser('ESP-IDF Partitions Tool') + quiet = quiet_flag - 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. - parser.add_argument( - '--port', - '-p', - help='port where the target device of the command is connected to; the partition table is sourced from ' - 'this device when the partition table file is not defined', - ) - parser.add_argument('--baud', '-b', help='baudrate to use', type=int) - - parser.add_argument('--partition-table-offset', '-o', help='offset to read the partition table from', type=str) - parser.add_argument('--primary-bootloader-offset', help='offset for primary bootloader', type=str) - parser.add_argument('--recovery-bootloader-offset', help='offset for recovery bootloader', type=str) - parser.add_argument( - '--partition-table-file', - '-f', - help='file (CSV/binary) to read the partition table from; ' - 'overrides device attached to specified port as the partition table source when defined', - ) - - partition_selection_parser = argparse.ArgumentParser(add_help=False) - - # Specify what partition to perform the operation on. This can either be specified using the - # partition name or the first partition that matches the specified type/subtype - partition_selection_args = partition_selection_parser.add_mutually_exclusive_group() - - partition_selection_args.add_argument('--partition-name', '-n', help='name of the partition') - partition_selection_args.add_argument('--partition-type', '-t', help='type of the partition') - partition_selection_args.add_argument( - '--partition-boot-default', - '-d', - help='select the default boot partition \ - using the same fallback logic as the IDF bootloader', - action='store_true', - ) - - partition_selection_parser.add_argument('--partition-subtype', '-s', help='subtype of the partition') - partition_selection_parser.add_argument( - '--extra-partition-subtypes', help='Extra partition subtype entries', nargs='*' - ) - - subparsers = parser.add_subparsers(dest='operation', help='run parttool -h for additional help') - - # Specify the supported operations - read_part_subparser = subparsers.add_parser( - 'read_partition', - help='read partition from device and dump contents into a file', - parents=[partition_selection_parser], - ) - read_part_subparser.add_argument('--output', help='file to dump the read partition contents to') - - write_part_subparser = subparsers.add_parser( - 'write_partition', - help='write contents of a binary file to partition on device', - parents=[partition_selection_parser], - ) - write_part_subparser.add_argument('--input', help='file whose contents are to be written to the partition offset') - write_part_subparser.add_argument('--ignore-readonly', help='Ignore read-only attribute', action='store_true') - - subparsers.add_parser( - 'erase_partition', help='erase the contents of a partition on the device', parents=[partition_selection_parser] - ) - - print_partition_info_subparser = subparsers.add_parser( - 'get_partition_info', help='get partition information', parents=[partition_selection_parser] - ) - print_partition_info_subparser.add_argument( - '--info', - help='type of partition information to get', - choices=['name', 'type', 'subtype', 'offset', 'size', 'encrypted', 'readonly'], - default=['offset', 'size'], - nargs='+', - ) - print_partition_info_subparser.add_argument( - '--part_list', help='Get a list of partitions suitable for a given type', action='store_true' - ) - - args = parser.parse_args() - quiet = args.quiet - - # No operation specified, display help and exit - if args.operation is None: - if not quiet: - parser.print_help() - sys.exit(1) - - # Prepare the partition to perform operation on - if args.partition_name: - partition_id = PartitionName(args.partition_name) - elif args.partition_type: - if not args.partition_subtype: - raise RuntimeError('--partition-subtype should be defined when --partition-type is defined') - partition_id = PartitionType(args.partition_type, args.partition_subtype, getattr(args, 'part_list', None)) - elif args.partition_boot_default: - partition_id = PARTITION_BOOT_DEFAULT - else: - raise RuntimeError( - 'Partition to operate on should be defined using --partition-name OR \ - partition-type,--partition-subtype OR partition-boot-default' - ) - - # Prepare the device to perform operation on - target_args = {} - - if args.port: - target_args['port'] = args.port - - if args.baud: - target_args['baud'] = args.baud - - if args.partition_table_file: - target_args['partition_table_file'] = args.partition_table_file - - if args.partition_table_offset: - target_args['partition_table_offset'] = int(args.partition_table_offset, 0) - - if args.primary_bootloader_offset: - target_args['primary_bootloader_offset'] = int(args.primary_bootloader_offset, 0) - - if args.recovery_bootloader_offset: - target_args['recovery_bootloader_offset'] = int(args.recovery_bootloader_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 - - if args.extra_partition_subtypes: - gen.add_extra_subtypes(args.extra_partition_subtypes) - - target = ParttoolTarget(**target_args) - - # Create the operation table and execute the operation - common_args = {'target': target, 'partition_id': partition_id} parttool_ops = { 'erase_partition': (_erase_partition, []), 'read_partition': (_read_partition, ['output']), @@ -385,13 +325,12 @@ def main(): 'get_partition_info': (_get_partition_info, ['info']), } - (op, op_args) = parttool_ops[args.operation] - - for op_arg in op_args: - common_args.update({op_arg: vars(args)[op_arg]}) + op, op_arg_names = parttool_ops[operation] + common_args = {'target': target, 'partition_id': partition_id} + for op_arg in op_arg_names: + common_args[op_arg] = op_kwargs[op_arg] if quiet: - # If exceptions occur, suppress and exit quietly try: op(**common_args) except Exception: @@ -400,9 +339,196 @@ def main(): try: op(**common_args) except gen.InputError as e: - print(e, file=sys.stderr) + log.print(str(e), file=sys.stderr) sys.exit(2) +@click.group( + cls=EspRichGroup, + context_settings={'help_option_names': ['-h', '--help']}, + help='ESP-IDF Partitions Tool', +) +@click.option('--quiet', '-q', is_flag=True, help='suppress stderr messages') +@click.option( + '--esptool-args', + cls=OptionEatAll, + type=_ESPTOOL_ARGS_TYPE, + default=(), + help='additional main arguments for esptool', +) +@click.option( + '--esptool-write-args', + cls=OptionEatAll, + type=_ESPTOOL_ARGS_TYPE, + default=(), + help='additional subcommand arguments when writing to flash', +) +@click.option( + '--esptool-read-args', + cls=OptionEatAll, + type=_ESPTOOL_ARGS_TYPE, + default=(), + help='additional subcommand arguments when reading flash', +) +@click.option( + '--esptool-erase-args', + cls=OptionEatAll, + type=_ESPTOOL_ARGS_TYPE, + default=(), + help='additional subcommand arguments when erasing regions of flash', +) +@click.option( + '--port', + '-p', + help='port where the target device of the command is connected to; the partition table is sourced from ' + 'this device when the partition table file is not defined', +) +@click.option('--baud', '-b', type=int, help='baudrate to use') +@click.option('--partition-table-offset', '-o', type=AnyIntType(), help='offset to read the partition table from') +@click.option('--primary-bootloader-offset', type=AnyIntType(), help='offset for primary bootloader') +@click.option('--recovery-bootloader-offset', type=AnyIntType(), help='offset for recovery bootloader') +@click.option( + '--partition-table-file', + '-f', + type=click.Path(), + help='file (CSV/binary) to read the partition table from; ' + 'overrides device attached to specified port as the partition table source when defined', +) +@click.pass_context +def cli( + ctx, + quiet, + esptool_args, + esptool_write_args, + esptool_read_args, + esptool_erase_args, + port, + baud, + partition_table_offset, + primary_bootloader_offset, + recovery_bootloader_offset, + partition_table_file, +): + ctx.ensure_object(dict) + ctx.obj.update( + { + 'quiet': quiet, + 'esptool_args': esptool_args, + 'esptool_write_args': esptool_write_args, + 'esptool_read_args': esptool_read_args, + 'esptool_erase_args': esptool_erase_args, + 'port': port, + 'baud': baud, + 'partition_table_offset': partition_table_offset, + 'primary_bootloader_offset': primary_bootloader_offset, + 'recovery_bootloader_offset': recovery_bootloader_offset, + 'partition_table_file': partition_table_file, + } + ) + + +def _shared_subcommand_options(func): + func = click.pass_context(func) + return func + + +@cli.command('read_partition', help='read partition from device and dump contents into a file') +@_partition_selection_options +@click.option('--output', help='file to dump the read partition contents to') +@_shared_subcommand_options +def read_partition_cmd( + ctx, + partition_name, + partition_type, + partition_boot_default, + partition_subtype, + extra_partition_subtypes, + output, +): + partition_id = _build_partition_id(partition_name, partition_type, partition_subtype, partition_boot_default) + target = ParttoolTarget(**_target_kwargs_from_ctx(ctx.obj, extra_partition_subtypes)) + _run_operation('read_partition', ctx.obj['quiet'], target, partition_id, output=output) + + +@cli.command('write_partition', help='write contents of a binary file to partition on device') +@_partition_selection_options +@click.option('--input', 'input_file', help='file whose contents are to be written to the partition offset') +@click.option('--ignore-readonly', is_flag=True, help='Ignore read-only attribute') +@_shared_subcommand_options +def write_partition_cmd( + ctx, + partition_name, + partition_type, + partition_boot_default, + partition_subtype, + extra_partition_subtypes, + input_file, + ignore_readonly, +): + partition_id = _build_partition_id(partition_name, partition_type, partition_subtype, partition_boot_default) + target = ParttoolTarget(**_target_kwargs_from_ctx(ctx.obj, extra_partition_subtypes)) + _run_operation( + 'write_partition', + ctx.obj['quiet'], + target, + partition_id, + input=input_file, + ignore_readonly=ignore_readonly, + ) + + +@cli.command('erase_partition', help='erase the contents of a partition on the device') +@_partition_selection_options +@_shared_subcommand_options +def erase_partition_cmd( + ctx, + partition_name, + partition_type, + partition_boot_default, + partition_subtype, + extra_partition_subtypes, +): + partition_id = _build_partition_id(partition_name, partition_type, partition_subtype, partition_boot_default) + target = ParttoolTarget(**_target_kwargs_from_ctx(ctx.obj, extra_partition_subtypes)) + _run_operation('erase_partition', ctx.obj['quiet'], target, partition_id) + + +@cli.command('get_partition_info', help='get partition information') +@_partition_selection_options +@click.option( + '--info', + type=click.Choice(['name', 'type', 'subtype', 'offset', 'size', 'encrypted', 'readonly']), + multiple=True, + cls=OptionEatAll, + default=['offset', 'size'], + show_default=True, + help='type of partition information to get', +) +@click.option('--part_list', 'part_list', is_flag=True, help='Get a list of partitions suitable for a given type') +@_shared_subcommand_options +def get_partition_info_cmd( + ctx, + partition_name, + partition_type, + partition_boot_default, + partition_subtype, + extra_partition_subtypes, + info, + part_list, +): + partition_id = _build_partition_id( + partition_name, partition_type, partition_subtype, partition_boot_default, part_list=part_list + ) + target = ParttoolTarget(**_target_kwargs_from_ctx(ctx.obj, extra_partition_subtypes)) + _run_operation('get_partition_info', ctx.obj['quiet'], target, partition_id, info=info) + + +def main(): + cli() + + if __name__ == '__main__': + from esp_pylib.excepthook import install_exception_reporting + + install_exception_reporting() main()