xref: /aosp_15_r20/external/pigweed/pw_rpc/py/pw_rpc/plugin.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"""pw_rpc protoc plugin entrypoint to generate code for RPC services."""
15
16import enum
17import sys
18
19from google.protobuf.compiler import plugin_pb2
20
21from pw_rpc import codegen_nanopb
22from pw_rpc import codegen_pwpb
23from pw_rpc import codegen_raw
24
25
26class Codegen(enum.Enum):
27    RAW = 0
28    NANOPB = 1
29    PWPB = 2
30
31
32def process_proto_request(
33    codegen: Codegen,
34    req: plugin_pb2.CodeGeneratorRequest,
35    res: plugin_pb2.CodeGeneratorResponse,
36) -> None:
37    """Handles a protoc CodeGeneratorRequest message.
38
39    Generates code for the files in the request and writes the output to the
40    specified CodeGeneratorResponse message.
41
42    Args:
43      req: A CodeGeneratorRequest for a proto compilation.
44      res: A CodeGeneratorResponse to populate with the plugin's output.
45    """
46    for proto_file in req.proto_file:
47        if codegen is Codegen.RAW:
48            output_files = codegen_raw.process_proto_file(proto_file)
49        elif codegen is Codegen.NANOPB:
50            output_files = codegen_nanopb.process_proto_file(proto_file)
51        elif codegen is Codegen.PWPB:
52            output_files = codegen_pwpb.process_proto_file(proto_file)
53        else:
54            raise NotImplementedError(f'Unknown codegen type {codegen}')
55
56        for output_file in output_files:
57            fd = res.file.add()
58            fd.name = output_file.name()
59            fd.content = output_file.content()
60
61
62def main(codegen: Codegen) -> int:
63    """Protobuf compiler plugin entrypoint.
64
65    Reads a CodeGeneratorRequest proto from stdin and writes a
66    CodeGeneratorResponse to stdout.
67    """
68    data = sys.stdin.buffer.read()
69    request = plugin_pb2.CodeGeneratorRequest.FromString(data)
70    response = plugin_pb2.CodeGeneratorResponse()
71    process_proto_request(codegen, request, response)
72
73    # Declare that this plugin supports optional fields in proto3. No proto
74    # message code is generated, so optional in proto3 is supported trivially.
75    response.supported_features |= (  # type: ignore[attr-defined]
76        response.FEATURE_PROTO3_OPTIONAL
77    )  # type: ignore[attr-defined]
78
79    sys.stdout.buffer.write(response.SerializeToString())
80    return 0
81