third_party/pigweed/repo/pw_rpc/py/pw_rpc/plugin.py

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