• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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