third_party/pigweed/repo/pw_bloat/py/bloat.py

   1 # Copyright 2019 The Pigweed Authors
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
   4 # use this file except in compliance with the License. You may obtain a copy of
   5 # the License at
   6 #
   7 #     https://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  11 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  12 # License for the specific language governing permissions and limitations under
  13 # the License.
  14 """
  15 bloat is a script which generates a size report card for binary files.
  16 """
  17
  18 import argparse
  19 import logging
  20 import os
  21 import subprocess
  22 import sys
  23
  24 from typing import List, Iterable, Optional
  25
  26 from binary_diff import BinaryDiff
  27 import bloat_output
  28
  29 import pw_cli.log
  30
  31 _LOG = logging.getLogger(__name__)
  32
  33
  34 def parse_args() -> argparse.Namespace:
  35     """Parses the script's arguments."""
  36     def delimited_list(delimiter: str, items: Optional[int] = None):
  37         def _parser(arg: str):
  38             args = arg.split(delimiter)
  39
  40             if items and len(args) != items:
  41                 raise argparse.ArgumentTypeError(
  42                     'Argument must be a '
  43                     f'{delimiter}-delimited list with {items} items: "{arg}"')
  44
  45             return args
  46
  47         return _parser
  48
  49     parser = argparse.ArgumentParser(
  50         'Generate a size report card for binaries')
  51     parser.add_argument('--bloaty-config',
  52                         type=delimited_list(';'),
  53                         required=True,
  54                         help='Data source configuration for Bloaty')
  55     parser.add_argument('--full',
  56                         action='store_true',
  57                         help='Display full bloat breakdown by symbol')
  58     parser.add_argument('--labels',
  59                         type=delimited_list(';'),
  60                         default='',
  61                         help='Labels for output binaries')
  62     parser.add_argument('--out-dir',
  63                         type=str,
  64                         required=True,
  65                         help='Directory in which to write output files')
  66     parser.add_argument('--target',
  67                         type=str,
  68                         required=True,
  69                         help='Build target name')
  70     parser.add_argument('--title',
  71                         type=str,
  72                         default='pw_bloat',
  73                         help='Report title')
  74     parser.add_argument('--source-filter',
  75                         type=str,
  76                         help='Bloaty data source filter')
  77     parser.add_argument('diff_targets',
  78                         type=delimited_list(';', 2),
  79                         nargs='+',
  80                         metavar='DIFF_TARGET',
  81                         help='Binary;base pairs to process')
  82
  83     return parser.parse_args()
  84
  85
  86 def run_bloaty(
  87     filename: str,
  88     config: str,
  89     base_file: Optional[str] = None,
  90     data_sources: Iterable[str] = (),
  91     extra_args: Iterable[str] = ()
  92 ) -> bytes:
  93     """Executes a Bloaty size report on some binary file(s).
  94
  95     Args:
  96         filename: Path to the binary.
  97         config: Path to Bloaty config file.
  98         base_file: Path to a base binary. If provided, a size diff is performed.
  99         data_sources: List of Bloaty data sources for the report.
 100         extra_args: Additional command-line arguments to pass to Bloaty.
 101
 102     Returns:
 103         Binary output of the Bloaty invocation.
 104
 105     Raises:
 106         subprocess.CalledProcessError: The Bloaty invocation failed.
 107     """
 108
 109     # TODO(frolv): Point the default bloaty path to a prebuilt in Pigweed.
 110     default_bloaty = 'bloaty'
 111     bloaty_path = os.getenv('BLOATY_PATH', default_bloaty)
 112
 113     # yapf: disable
 114     cmd = [
 115         bloaty_path,
 116         '-c', config,
 117         '-d', ','.join(data_sources),
 118         '--domain', 'vm',
 119         filename,
 120         *extra_args
 121     ]
 122     # yapf: enable
 123
 124     if base_file is not None:
 125         cmd.extend(['--', base_file])
 126
 127     return subprocess.check_output(cmd)
 128
 129
 130 def main() -> int:
 131     """Program entry point."""
 132
 133     args = parse_args()
 134
 135     base_binaries: List[str] = []
 136     diff_binaries: List[str] = []
 137
 138     try:
 139         for binary, base in args.diff_targets:
 140             diff_binaries.append(binary)
 141             base_binaries.append(base)
 142     except RuntimeError as err:
 143         _LOG.error('%s: %s', sys.argv[0], err)
 144         return 1
 145
 146     data_sources = ['segment_names']
 147     if args.full:
 148         data_sources.append('fullsymbols')
 149
 150     # TODO(frolv): CSV output is disabled for full reports as the default Bloaty
 151     # breakdown is printed. This script should be modified to print a custom
 152     # symbol breakdown in full reports.
 153     extra_args = [] if args.full else ['--csv']
 154     if args.source_filter:
 155         extra_args.extend(['--source-filter', args.source_filter])
 156
 157     diffs: List[BinaryDiff] = []
 158     report = []
 159
 160     for i, binary in enumerate(diff_binaries):
 161         binary_name = (args.labels[i]
 162                        if i < len(args.labels) else os.path.basename(binary))
 163         try:
 164             output = run_bloaty(binary, args.bloaty_config[i],
 165                                 base_binaries[i], data_sources, extra_args)
 166             if not output:
 167                 continue
 168
 169             # TODO(frolv): Remove when custom output for full mode is added.
 170             if args.full:
 171                 report.append(binary_name)
 172                 report.append('-' * len(binary_name))
 173                 report.append(output.decode())
 174                 continue
 175
 176             # Ignore the first row as it displays column names.
 177             bloaty_csv = output.decode().splitlines()[1:]
 178             diffs.append(BinaryDiff.from_csv(binary_name, bloaty_csv))
 179         except subprocess.CalledProcessError:
 180             _LOG.error('%s: failed to run diff on %s', sys.argv[0], binary)
 181             return 1
 182
 183     def write_file(filename: str, contents: str) -> None:
 184         path = os.path.join(args.out_dir, filename)
 185         with open(path, 'w') as output_file:
 186             output_file.write(contents)
 187         _LOG.debug('Output written to %s', path)
 188
 189     # TODO(frolv): Remove when custom output for full mode is added.
 190     if not args.full:
 191         out = bloat_output.TableOutput(args.title,
 192                                        diffs,
 193                                        charset=bloat_output.LineCharset)
 194         report.append(out.diff())
 195
 196         rst = bloat_output.RstOutput(diffs)
 197         write_file(f'{args.target}', rst.diff())
 198
 199     complete_output = '\n'.join(report) + '\n'
 200     write_file(f'{args.target}.txt', complete_output)
 201     print(complete_output)
 202
 203     return 0
 204
 205
 206 if __name__ == '__main__':
 207     pw_cli.log.install()
 208     sys.exit(main())