third_party/pigweed/repo/pw_protobuf_compiler/py/pw_protobuf_compiler/generate_protos.py

   1 # Copyright 2020 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 """Script that invokes protoc to generate code for .proto files."""
  15
  16 import argparse
  17 import logging
  18 import os
  19 from pathlib import Path
  20 import subprocess
  21 import sys
  22 import tempfile
  23
  24 from typing import Callable, Dict, Optional, Tuple
  25
  26 # Make sure dependencies are optional, since this script may be run when
  27 # installing Python package dependencies through GN.
  28 try:
  29     from pw_cli.log import install as setup_logging
  30 except ImportError:
  31     from logging import basicConfig as setup_logging  # type: ignore
  32
  33 _LOG = logging.getLogger(__name__)
  34
  35 _COMMON_FLAGS = ('--experimental_allow_proto3_optional', )
  36
  37
  38 def argument_parser(
  39     parser: Optional[argparse.ArgumentParser] = None
  40 ) -> argparse.ArgumentParser:
  41     """Registers the script's arguments on an argument parser."""
  42
  43     if parser is None:
  44         parser = argparse.ArgumentParser(description=__doc__)
  45
  46     parser.add_argument('--language',
  47                         required=True,
  48                         choices=DEFAULT_PROTOC_ARGS,
  49                         help='Output language')
  50     parser.add_argument('--plugin-path',
  51                         type=Path,
  52                         help='Path to the protoc plugin')
  53     parser.add_argument('--include-path',
  54                         required=True,
  55                         help='Include path for proto compilation')
  56     parser.add_argument('--include-file',
  57                         type=argparse.FileType('r'),
  58                         help='File containing additional protoc include paths')
  59     parser.add_argument('--out-dir',
  60                         required=True,
  61                         help='Output directory for generated code')
  62     parser.add_argument('protos',
  63                         metavar='PROTO',
  64                         nargs='+',
  65                         help='Input protobuf files')
  66
  67     return parser
  68
  69
  70 def protoc_cc_args(args: argparse.Namespace) -> Tuple[str, ...]:
  71     return _COMMON_FLAGS + (
  72         '--plugin',
  73         f'protoc-gen-custom={args.plugin_path}',
  74         '--custom_out',
  75         args.out_dir,
  76     )
  77
  78
  79 def protoc_go_args(args: argparse.Namespace) -> Tuple[str, ...]:
  80     return _COMMON_FLAGS + (
  81         '--go_out',
  82         f'plugins=grpc:{args.out_dir}',
  83     )
  84
  85
  86 def protoc_nanopb_args(args: argparse.Namespace) -> Tuple[str, ...]:
  87     # nanopb needs to know of the include path to parse *.options files
  88     return _COMMON_FLAGS + (
  89         '--plugin',
  90         f'protoc-gen-nanopb={args.plugin_path}',
  91         # nanopb_opt provides the flags to use for nanopb_out. Windows doesn't
  92         # like when you merge the two using the `flag,...:out` syntax.
  93         f'--nanopb_opt=-I{args.include_path}',
  94         f'--nanopb_out={args.out_dir}',
  95     )
  96
  97
  98 def protoc_nanopb_rpc_args(args: argparse.Namespace) -> Tuple[str, ...]:
  99     return _COMMON_FLAGS + (
 100         '--plugin',
 101         f'protoc-gen-custom={args.plugin_path}',
 102         '--custom_out',
 103         args.out_dir,
 104     )
 105
 106
 107 def protoc_raw_rpc_args(args: argparse.Namespace) -> Tuple[str, ...]:
 108     return _COMMON_FLAGS + (
 109         '--plugin',
 110         f'protoc-gen-custom={args.plugin_path}',
 111         '--custom_out',
 112         args.out_dir,
 113     )
 114
 115
 116 def protoc_python_args(args: argparse.Namespace) -> Tuple[str, ...]:
 117     return _COMMON_FLAGS + (
 118         '--python_out',
 119         args.out_dir,
 120         '--mypy_out',
 121         args.out_dir,
 122     )
 123
 124
 125 _DefaultArgsFunction = Callable[[argparse.Namespace], Tuple[str, ...]]
 126
 127 # Default additional protoc arguments for each supported language.
 128 # TODO(frolv): Make these overridable with a command-line argument.
 129 DEFAULT_PROTOC_ARGS: Dict[str, _DefaultArgsFunction] = {
 130     'pwpb': protoc_cc_args,
 131     'go': protoc_go_args,
 132     'nanopb': protoc_nanopb_args,
 133     'nanopb_rpc': protoc_nanopb_rpc_args,
 134     'raw_rpc': protoc_raw_rpc_args,
 135     'python': protoc_python_args,
 136 }
 137
 138 # Languages that protoc internally supports.
 139 BUILTIN_PROTOC_LANGS = ('go', 'python')
 140
 141
 142 def main() -> int:
 143     """Runs protoc as configured by command-line arguments."""
 144
 145     parser = argument_parser()
 146     args = parser.parse_args()
 147
 148     if args.plugin_path is None and args.language not in BUILTIN_PROTOC_LANGS:
 149         parser.error(
 150             f'--plugin-path is required for --language {args.language}')
 151
 152     os.makedirs(args.out_dir, exist_ok=True)
 153
 154     include_paths = [f'-I{line.strip()}' for line in args.include_file]
 155
 156     wrapper_script: Optional[Path] = None
 157
 158     # On Windows, use a .bat version of the plugin if it exists or create a .bat
 159     # wrapper to use if none exists.
 160     if os.name == 'nt' and args.plugin_path:
 161         if args.plugin_path.with_suffix('.bat').exists():
 162             args.plugin_path = args.plugin_path.with_suffix('.bat')
 163             _LOG.debug('Using Batch plugin %s', args.plugin_path)
 164         else:
 165             with tempfile.NamedTemporaryFile('w', suffix='.bat',
 166                                              delete=False) as file:
 167                 file.write(f'@echo off\npython {args.plugin_path.resolve()}\n')
 168
 169             args.plugin_path = wrapper_script = Path(file.name)
 170             _LOG.debug('Using generated plugin wrapper %s', args.plugin_path)
 171
 172     try:
 173         process = subprocess.run(
 174             [
 175                 'protoc',
 176                 f'-I{args.include_path}',
 177                 *include_paths,
 178                 *DEFAULT_PROTOC_ARGS[args.language](args),
 179                 *args.protos,
 180             ],
 181             stdout=subprocess.PIPE,
 182             stderr=subprocess.STDOUT,
 183         )
 184     finally:
 185         if wrapper_script:
 186             wrapper_script.unlink()
 187
 188     if process.returncode != 0:
 189         sys.stderr.buffer.write(process.stdout)
 190         sys.stderr.flush()
 191
 192     return process.returncode
 193
 194
 195 if __name__ == '__main__':
 196     setup_logging()
 197     sys.exit(main())