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 16import argparse 17import logging 18import os 19from pathlib import Path 20import subprocess 21import sys 22import tempfile 23from typing import Callable, Dict, List, Optional, Tuple, Union 24 25# Make sure dependencies are optional, since this script may be run when 26# installing Python package dependencies through GN. 27try: 28 from pw_cli.log import install as setup_logging 29except ImportError: 30 from logging import basicConfig as setup_logging # type: ignore 31 32_LOG = logging.getLogger(__name__) 33 34 35def _argument_parser() -> argparse.ArgumentParser: 36 """Registers the script's arguments on an argument parser.""" 37 38 parser = argparse.ArgumentParser(description=__doc__) 39 40 parser.add_argument( 41 '--language', 42 required=True, 43 choices=DEFAULT_PROTOC_ARGS, 44 help='Output language', 45 ) 46 parser.add_argument( 47 '--plugin-path', type=Path, help='Path to the protoc plugin' 48 ) 49 parser.add_argument( 50 '--proto-path', 51 type=Path, 52 help='Additional protoc include paths', 53 action='append', 54 ) 55 parser.add_argument( 56 '--include-file', 57 type=argparse.FileType('r'), 58 help='File containing additional protoc include paths', 59 ) 60 parser.add_argument( 61 '--out-dir', 62 type=Path, 63 required=True, 64 help='Output directory for generated code', 65 ) 66 parser.add_argument( 67 '--compile-dir', 68 type=Path, 69 required=True, 70 help='Root path for compilation', 71 ) 72 parser.add_argument( 73 '--sources', type=Path, nargs='+', help='Input protobuf files' 74 ) 75 parser.add_argument( 76 '--protoc', type=Path, default='protoc', help='Path to protoc' 77 ) 78 parser.add_argument( 79 '--no-experimental-proto3-optional', 80 dest='experimental_proto3_optional', 81 action='store_false', 82 help='Do not invoke protoc with --experimental_allow_proto3_optional', 83 ) 84 parser.add_argument( 85 '--no-generate-type-hints', 86 dest='generate_type_hints', 87 action='store_false', 88 help='Do not generate pyi files for python', 89 ) 90 parser.add_argument( 91 '--exclude-pwpb-legacy-snake-case-field-name-enums', 92 dest='exclude_pwpb_legacy_snake_case_field_name_enums', 93 action='store_true', 94 help=( 95 'If set, generates legacy SNAKE_CASE names for field name enums ' 96 'in PWPB.' 97 ), 98 ) 99 100 return parser 101 102 103def protoc_common_args(args: argparse.Namespace) -> Tuple[str, ...]: 104 flags: Tuple[str, ...] = () 105 if args.experimental_proto3_optional: 106 flags += ('--experimental_allow_proto3_optional',) 107 return flags 108 109 110def protoc_pwpb_args( 111 args: argparse.Namespace, include_paths: List[str] 112) -> Tuple[str, ...]: 113 out_args = [ 114 '--plugin', 115 f'protoc-gen-custom={args.plugin_path}', 116 f'--custom_opt=-I{args.compile_dir}', 117 *[f'--custom_opt=-I{include_path}' for include_path in include_paths], 118 ] 119 120 if args.exclude_pwpb_legacy_snake_case_field_name_enums: 121 out_args.append( 122 '--custom_opt=--exclude-legacy-snake-case-field-name-enums' 123 ) 124 125 out_args.extend( 126 [ 127 '--custom_out', 128 args.out_dir, 129 ] 130 ) 131 132 return tuple(out_args) 133 134 135def protoc_pwpb_rpc_args( 136 args: argparse.Namespace, _include_paths: List[str] 137) -> Tuple[str, ...]: 138 return ( 139 '--plugin', 140 f'protoc-gen-custom={args.plugin_path}', 141 '--custom_out', 142 args.out_dir, 143 ) 144 145 146def protoc_go_args( 147 args: argparse.Namespace, _include_paths: List[str] 148) -> Tuple[str, ...]: 149 return ( 150 '--go_out', 151 f'plugins=grpc:{args.out_dir}', 152 ) 153 154 155def protoc_nanopb_args( 156 args: argparse.Namespace, _include_paths: List[str] 157) -> Tuple[str, ...]: 158 # nanopb needs to know of the include path to parse *.options files 159 return ( 160 '--plugin', 161 f'protoc-gen-nanopb={args.plugin_path}', 162 # nanopb_opt provides the flags to use for nanopb_out. Windows doesn't 163 # like when you merge the two using the `flag,...:out` syntax. Use 164 # Posix-style paths since backslashes on Windows are treated like 165 # escape characters. 166 f'--nanopb_opt=-I{args.compile_dir.as_posix()}', 167 f'--nanopb_out={args.out_dir}', 168 ) 169 170 171def protoc_nanopb_rpc_args( 172 args: argparse.Namespace, _include_paths: List[str] 173) -> Tuple[str, ...]: 174 return ( 175 '--plugin', 176 f'protoc-gen-custom={args.plugin_path}', 177 '--custom_out', 178 args.out_dir, 179 ) 180 181 182def protoc_raw_rpc_args( 183 args: argparse.Namespace, _include_paths: List[str] 184) -> Tuple[str, ...]: 185 return ( 186 '--plugin', 187 f'protoc-gen-custom={args.plugin_path}', 188 '--custom_out', 189 args.out_dir, 190 ) 191 192 193def protoc_python_args( 194 args: argparse.Namespace, _include_paths: List[str] 195) -> Tuple[str, ...]: 196 flags: Tuple[str, ...] = ( 197 '--python_out', 198 args.out_dir, 199 ) 200 201 if args.generate_type_hints: 202 flags += ( 203 '--mypy_out', 204 args.out_dir, 205 ) 206 207 return flags 208 209 210_DefaultArgsFunction = Callable[ 211 [argparse.Namespace, List[str]], Tuple[str, ...] 212] 213 214# Default additional protoc arguments for each supported language. 215# TODO(frolv): Make these overridable with a command-line argument. 216DEFAULT_PROTOC_ARGS: Dict[str, _DefaultArgsFunction] = { 217 'go': protoc_go_args, 218 'nanopb': protoc_nanopb_args, 219 'nanopb_rpc': protoc_nanopb_rpc_args, 220 'pwpb': protoc_pwpb_args, 221 'pwpb_rpc': protoc_pwpb_rpc_args, 222 'python': protoc_python_args, 223 'raw_rpc': protoc_raw_rpc_args, 224} 225 226# Languages that protoc internally supports. 227BUILTIN_PROTOC_LANGS = ('go', 'python') 228 229 230def main() -> int: 231 """Runs protoc as configured by command-line arguments.""" 232 233 parser = _argument_parser() 234 args = parser.parse_args() 235 236 if args.plugin_path is None and args.language not in BUILTIN_PROTOC_LANGS: 237 parser.error( 238 f'--plugin-path is required for --language {args.language}' 239 ) 240 241 args.out_dir.mkdir(parents=True, exist_ok=True) 242 243 include_paths: List[str] = [] 244 if args.include_file: 245 include_paths.extend(line.strip() for line in args.include_file) 246 if args.proto_path: 247 include_paths.extend(str(path) for path in args.proto_path) 248 249 wrapper_script: Optional[Path] = None 250 251 # On Windows, use a .bat version of the plugin if it exists or create a .bat 252 # wrapper to use if none exists. 253 if os.name == 'nt' and args.plugin_path: 254 if args.plugin_path.with_suffix('.bat').exists(): 255 args.plugin_path = args.plugin_path.with_suffix('.bat') 256 _LOG.debug('Using Batch plugin %s', args.plugin_path) 257 else: 258 with tempfile.NamedTemporaryFile( 259 'w', suffix='.bat', delete=False 260 ) as file: 261 file.write(f'@echo off\npython {args.plugin_path.resolve()}\n') 262 263 args.plugin_path = wrapper_script = Path(file.name) 264 _LOG.debug('Using generated plugin wrapper %s', args.plugin_path) 265 266 cmd: Tuple[Union[str, Path], ...] = ( 267 args.protoc, 268 f'-I{args.compile_dir}', 269 *[f'-I{include_path}' for include_path in include_paths], 270 *protoc_common_args(args), 271 *DEFAULT_PROTOC_ARGS[args.language](args, include_paths), 272 *args.sources, 273 ) 274 275 try: 276 process = subprocess.run( 277 cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT 278 ) 279 finally: 280 if wrapper_script: 281 wrapper_script.unlink() 282 283 if process.returncode != 0: 284 _LOG.error( 285 'Protocol buffer compilation failed!\n%s', 286 ' '.join(str(c) for c in cmd), 287 ) 288 sys.stderr.buffer.write(process.stdout) 289 sys.stderr.flush() 290 291 return process.returncode 292 293 294if __name__ == '__main__': 295 setup_logging() 296 sys.exit(main()) 297