xref: /aosp_15_r20/external/pigweed/pw_protobuf_compiler/py/pw_protobuf_compiler/generate_protos.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
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
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-experimental-editions',
86        dest='experimental_editions',
87        action='store_false',
88        help='Do not invoke protoc with --experimental_editions',
89    )
90    parser.add_argument(
91        '--no-generate-type-hints',
92        dest='generate_type_hints',
93        action='store_false',
94        help='Do not generate pyi files for python',
95    )
96    parser.add_argument(
97        '--exclude-pwpb-legacy-snake-case-field-name-enums',
98        dest='exclude_pwpb_legacy_snake_case_field_name_enums',
99        action='store_true',
100        help=(
101            'If set, generates legacy SNAKE_CASE names for field name enums '
102            'in PWPB.'
103        ),
104    )
105    parser.add_argument(
106        '--pwpb-no-generic-options-files',
107        action='store_true',
108        help=(
109            'If set, requires the use of the `.pwpb_options` for pw_protobuf '
110            'options files'
111        ),
112    )
113    parser.add_argument(
114        '--pwpb-no-oneof-callbacks',
115        action='store_true',
116        help='Generate legacy inline oneof members instead of callbacks',
117    )
118
119    return parser
120
121
122def protoc_common_args(args: argparse.Namespace) -> tuple[str, ...]:
123    flags: tuple[str, ...] = ()
124    if args.experimental_proto3_optional:
125        flags += ('--experimental_allow_proto3_optional',)
126    if args.experimental_editions:
127        flags += ('--experimental_editions',)
128    return flags
129
130
131def protoc_pwpb_args(
132    args: argparse.Namespace, include_paths: list[str]
133) -> tuple[str, ...]:
134    out_args = [
135        '--plugin',
136        f'protoc-gen-custom={args.plugin_path}',
137        f'--custom_opt=-I{args.compile_dir}',
138        *[f'--custom_opt=-I{include_path}' for include_path in include_paths],
139    ]
140
141    if args.exclude_pwpb_legacy_snake_case_field_name_enums:
142        out_args.append(
143            '--custom_opt=--exclude-legacy-snake-case-field-name-enums'
144        )
145    if args.pwpb_no_generic_options_files:
146        out_args.append('--custom_opt=--no-generic-options-files')
147    if args.pwpb_no_oneof_callbacks:
148        out_args.append('--custom_opt=--no-oneof-callbacks')
149
150    out_args.extend(
151        [
152            '--custom_out',
153            args.out_dir,
154        ]
155    )
156
157    return tuple(out_args)
158
159
160def protoc_pwpb_rpc_args(
161    args: argparse.Namespace, _include_paths: list[str]
162) -> tuple[str, ...]:
163    return (
164        '--plugin',
165        f'protoc-gen-custom={args.plugin_path}',
166        '--custom_out',
167        args.out_dir,
168    )
169
170
171def protoc_go_args(
172    args: argparse.Namespace, _include_paths: list[str]
173) -> tuple[str, ...]:
174    return (
175        '--go_out',
176        f'plugins=grpc:{args.out_dir}',
177    )
178
179
180def protoc_nanopb_args(
181    args: argparse.Namespace, _include_paths: list[str]
182) -> tuple[str, ...]:
183    # nanopb needs to know of the include path to parse *.options files
184    return (
185        '--plugin',
186        f'protoc-gen-nanopb={args.plugin_path}',
187        # nanopb_opt provides the flags to use for nanopb_out. Windows doesn't
188        # like when you merge the two using the `flag,...:out` syntax. Use
189        # Posix-style paths since backslashes on Windows are treated like
190        # escape characters.
191        f'--nanopb_opt=-I{args.compile_dir.as_posix()}',
192        f'--nanopb_out={args.out_dir}',
193    )
194
195
196def protoc_nanopb_rpc_args(
197    args: argparse.Namespace, _include_paths: list[str]
198) -> tuple[str, ...]:
199    return (
200        '--plugin',
201        f'protoc-gen-custom={args.plugin_path}',
202        '--custom_out',
203        args.out_dir,
204    )
205
206
207def protoc_raw_rpc_args(
208    args: argparse.Namespace, _include_paths: list[str]
209) -> tuple[str, ...]:
210    return (
211        '--plugin',
212        f'protoc-gen-custom={args.plugin_path}',
213        '--custom_out',
214        args.out_dir,
215    )
216
217
218def protoc_python_args(
219    args: argparse.Namespace, _include_paths: list[str]
220) -> tuple[str, ...]:
221    flags: tuple[str, ...] = (
222        '--python_out',
223        args.out_dir,
224    )
225
226    if args.generate_type_hints:
227        flags += (
228            '--mypy_out',
229            args.out_dir,
230        )
231
232    return flags
233
234
235_DefaultArgsFunction = Callable[
236    [argparse.Namespace, list[str]], tuple[str, ...]
237]
238
239# Default additional protoc arguments for each supported language.
240# TODO(frolv): Make these overridable with a command-line argument.
241DEFAULT_PROTOC_ARGS: dict[str, _DefaultArgsFunction] = {
242    'go': protoc_go_args,
243    'nanopb': protoc_nanopb_args,
244    'nanopb_rpc': protoc_nanopb_rpc_args,
245    'pwpb': protoc_pwpb_args,
246    'pwpb_rpc': protoc_pwpb_rpc_args,
247    'python': protoc_python_args,
248    'raw_rpc': protoc_raw_rpc_args,
249}
250
251# Languages that protoc internally supports.
252BUILTIN_PROTOC_LANGS = ('go', 'python')
253
254
255def main() -> int:
256    """Runs protoc as configured by command-line arguments."""
257
258    parser = _argument_parser()
259    args = parser.parse_args()
260
261    if args.plugin_path is None and args.language not in BUILTIN_PROTOC_LANGS:
262        parser.error(
263            f'--plugin-path is required for --language {args.language}'
264        )
265
266    args.out_dir.mkdir(parents=True, exist_ok=True)
267
268    include_paths: list[str] = []
269    if args.include_file:
270        include_paths.extend(line.strip() for line in args.include_file)
271    if args.proto_path:
272        include_paths.extend(str(path) for path in args.proto_path)
273
274    wrapper_script: Path | None = None
275
276    # On Windows, use a .bat version of the plugin if it exists or create a .bat
277    # wrapper to use if none exists.
278    if os.name == 'nt' and args.plugin_path:
279        if args.plugin_path.with_suffix('.bat').exists():
280            args.plugin_path = args.plugin_path.with_suffix('.bat')
281            _LOG.debug('Using Batch plugin %s', args.plugin_path)
282        else:
283            with tempfile.NamedTemporaryFile(
284                'w', suffix='.bat', delete=False
285            ) as file:
286                file.write(f'@echo off\npython {args.plugin_path.resolve()}\n')
287
288            args.plugin_path = wrapper_script = Path(file.name)
289            _LOG.debug('Using generated plugin wrapper %s', args.plugin_path)
290
291    cmd: tuple[str | Path, ...] = (
292        args.protoc,
293        f'-I{args.compile_dir}',
294        *[f'-I{include_path}' for include_path in include_paths],
295        *protoc_common_args(args),
296        *DEFAULT_PROTOC_ARGS[args.language](args, include_paths),
297        *args.sources,
298    )
299
300    try:
301        process = subprocess.run(
302            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
303        )
304    finally:
305        if wrapper_script:
306            wrapper_script.unlink()
307
308    if process.returncode != 0:
309        _LOG.error(
310            'Protocol buffer compilation failed!\n%s',
311            ' '.join(str(c) for c in cmd),
312        )
313        sys.stderr.buffer.write(process.stdout)
314        sys.stderr.flush()
315
316    return process.returncode
317
318
319if __name__ == '__main__':
320    setup_logging()
321    sys.exit(main())
322