fix(partition_table): normalize CLI output with esp-pylib

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()