• 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_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: List[str] = []
155    if args.include_file:
156        include_paths = [f'-I{line.strip()}' for line in args.include_file]
157
158    wrapper_script: Optional[Path] = None
159
160    # On Windows, use a .bat version of the plugin if it exists or create a .bat
161    # wrapper to use if none exists.
162    if os.name == 'nt' and args.plugin_path:
163        if args.plugin_path.with_suffix('.bat').exists():
164            args.plugin_path = args.plugin_path.with_suffix('.bat')
165            _LOG.debug('Using Batch plugin %s', args.plugin_path)
166        else:
167            with tempfile.NamedTemporaryFile('w', suffix='.bat',
168                                             delete=False) as file:
169                file.write(f'@echo off\npython {args.plugin_path.resolve()}\n')
170
171            args.plugin_path = wrapper_script = Path(file.name)
172            _LOG.debug('Using generated plugin wrapper %s', args.plugin_path)
173
174    cmd: Tuple[Union[str, Path], ...] = (
175        'protoc',
176        f'-I{args.compile_dir}',
177        *include_paths,
178        *DEFAULT_PROTOC_ARGS[args.language](args),
179        *args.sources,
180    )
181
182    try:
183        process = subprocess.run(cmd,
184                                 stdout=subprocess.PIPE,
185                                 stderr=subprocess.STDOUT)
186    finally:
187        if wrapper_script:
188            wrapper_script.unlink()
189
190    if process.returncode != 0:
191        _LOG.error('Protocol buffer compilation failed!\n%s',
192                   ' '.join(str(c) for c in cmd))
193        sys.stderr.buffer.write(process.stdout)
194        sys.stderr.flush()
195
196    return process.returncode
197
198
199if __name__ == '__main__':
200    setup_logging()
201    sys.exit(main())
202