• 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, 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_COMMON_FLAGS = ('--experimental_allow_proto3_optional', )
35
36
37def _argument_parser() -> argparse.ArgumentParser:
38    """Registers the script's arguments on an argument parser."""
39
40    parser = argparse.ArgumentParser(description=__doc__)
41
42    parser.add_argument('--language',
43                        required=True,
44                        choices=DEFAULT_PROTOC_ARGS,
45                        help='Output language')
46    parser.add_argument('--plugin-path',
47                        type=Path,
48                        help='Path to the protoc plugin')
49    parser.add_argument('--include-file',
50                        type=argparse.FileType('r'),
51                        help='File containing additional protoc include paths')
52    parser.add_argument('--out-dir',
53                        type=Path,
54                        required=True,
55                        help='Output directory for generated code')
56    parser.add_argument('--compile-dir',
57                        type=Path,
58                        required=True,
59                        help='Root path for compilation')
60    parser.add_argument('--sources',
61                        type=Path,
62                        nargs='+',
63                        help='Input protobuf files')
64
65    return parser
66
67
68def protoc_cc_args(args: argparse.Namespace) -> Tuple[str, ...]:
69    return _COMMON_FLAGS + (
70        '--plugin',
71        f'protoc-gen-custom={args.plugin_path}',
72        '--custom_out',
73        args.out_dir,
74    )
75
76
77def protoc_go_args(args: argparse.Namespace) -> Tuple[str, ...]:
78    return _COMMON_FLAGS + (
79        '--go_out',
80        f'plugins=grpc:{args.out_dir}',
81    )
82
83
84def protoc_nanopb_args(args: argparse.Namespace) -> Tuple[str, ...]:
85    # nanopb needs to know of the include path to parse *.options files
86    return _COMMON_FLAGS + (
87        '--plugin',
88        f'protoc-gen-nanopb={args.plugin_path}',
89        # nanopb_opt provides the flags to use for nanopb_out. Windows doesn't
90        # like when you merge the two using the `flag,...:out` syntax. Use
91        # Posix-style paths since backslashes on Windows are treated like
92        # escape characters.
93        f'--nanopb_opt=-I{args.compile_dir.as_posix()}',
94        f'--nanopb_out={args.out_dir}',
95    )
96
97
98def 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
107def 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
116def 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.
129DEFAULT_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.
139BUILTIN_PROTOC_LANGS = ('go', 'python')
140
141
142def 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    args.out_dir.mkdir(parents=True, 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    cmd: Tuple[Union[str, Path], ...] = (
173        'protoc',
174        f'-I{args.compile_dir}',
175        *include_paths,
176        *DEFAULT_PROTOC_ARGS[args.language](args),
177        *args.sources,
178    )
179
180    try:
181        process = subprocess.run(cmd,
182                                 stdout=subprocess.PIPE,
183                                 stderr=subprocess.STDOUT)
184    finally:
185        if wrapper_script:
186            wrapper_script.unlink()
187
188    if process.returncode != 0:
189        _LOG.error('Protocol buffer compilation failed!\n%s',
190                   ' '.join(str(c) for c in cmd))
191        sys.stderr.buffer.write(process.stdout)
192        sys.stderr.flush()
193
194    return process.returncode
195
196
197if __name__ == '__main__':
198    setup_logging()
199    sys.exit(main())
200