third_party/pigweed/repo/pw_rpc/py/pw_rpc/codegen.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 """Common RPC codegen utilities."""
  15
  16 from datetime import datetime
  17 import os
  18 from typing import cast, Any, Callable, Iterable
  19
  20 from pw_protobuf.output_file import OutputFile
  21 from pw_protobuf.proto_tree import ProtoNode, ProtoService, ProtoServiceMethod
  22
  23 import pw_rpc.ids
  24
  25 PLUGIN_NAME = 'pw_rpc_codegen'
  26 PLUGIN_VERSION = '0.2.0'
  27
  28 RPC_NAMESPACE = '::pw::rpc'
  29
  30 STUB_REQUEST_TODO = (
  31     '// TODO: Read the request as appropriate for your application')
  32 STUB_RESPONSE_TODO = (
  33     '// TODO: Fill in the response as appropriate for your application')
  34 STUB_WRITER_TODO = (
  35     '// TODO: Send responses with the writer as appropriate for your '
  36     'application')
  37
  38 ServerWriterGenerator = Callable[[OutputFile], None]
  39 MethodGenerator = Callable[[ProtoServiceMethod, int, OutputFile], None]
  40 ServiceGenerator = Callable[[ProtoService, ProtoNode, OutputFile], None]
  41 IncludesGenerator = Callable[[Any, ProtoNode], Iterable[str]]
  42
  43
  44 def package(file_descriptor_proto, proto_package: ProtoNode,
  45             output: OutputFile, includes: IncludesGenerator,
  46             service: ServiceGenerator, client: ServiceGenerator) -> None:
  47     """Generates service and client code for a package."""
  48     assert proto_package.type() == ProtoNode.Type.PACKAGE
  49
  50     output.write_line(f'// {os.path.basename(output.name())} automatically '
  51                       f'generated by {PLUGIN_NAME} {PLUGIN_VERSION}')
  52     output.write_line(f'// on {datetime.now().isoformat()}')
  53     output.write_line('// clang-format off')
  54     output.write_line('#pragma once\n')
  55
  56     output.write_line('#include <array>')
  57     output.write_line('#include <cstdint>')
  58     output.write_line('#include <type_traits>\n')
  59
  60     include_lines = [
  61         '#include "pw_rpc/internal/method_lookup.h"',
  62         '#include "pw_rpc/server_context.h"',
  63         '#include "pw_rpc/service.h"',
  64     ]
  65     include_lines += includes(file_descriptor_proto, proto_package)
  66
  67     for include_line in sorted(include_lines):
  68         output.write_line(include_line)
  69
  70     output.write_line()
  71
  72     if proto_package.cpp_namespace():
  73         file_namespace = proto_package.cpp_namespace()
  74         if file_namespace.startswith('::'):
  75             file_namespace = file_namespace[2:]
  76
  77         output.write_line(f'namespace {file_namespace} {{')
  78
  79     for node in proto_package:
  80         if node.type() == ProtoNode.Type.SERVICE:
  81             service(cast(ProtoService, node), proto_package, output)
  82             client(cast(ProtoService, node), proto_package, output)
  83
  84     if proto_package.cpp_namespace():
  85         output.write_line(f'}}  // namespace {file_namespace}')
  86
  87
  88 def service_class(service: ProtoService, root: ProtoNode, output: OutputFile,
  89                   server_writer_alias: ServerWriterGenerator,
  90                   method_union: str,
  91                   method_descriptor: MethodGenerator) -> None:
  92     """Generates a C++ derived class for a nanopb RPC service."""
  93
  94     output.write_line('namespace generated {')
  95
  96     base_class = f'{RPC_NAMESPACE}::Service'
  97     output.write_line('\ntemplate <typename Implementation>')
  98     output.write_line(
  99         f'class {service.cpp_namespace(root)} : public {base_class} {{')
 100     output.write_line(' public:')
 101
 102     with output.indent():
 103         output.write_line(
 104             f'using ServerContext = {RPC_NAMESPACE}::ServerContext;')
 105         server_writer_alias(output)
 106         output.write_line()
 107
 108         output.write_line(f'constexpr {service.name()}()')
 109         output.write_line(f'    : {base_class}(kServiceId, kMethods) {{}}')
 110
 111         output.write_line()
 112         output.write_line(
 113             f'{service.name()}(const {service.name()}&) = delete;')
 114         output.write_line(f'{service.name()}& operator='
 115                           f'(const {service.name()}&) = delete;')
 116
 117         output.write_line()
 118         output.write_line(f'static constexpr const char* name() '
 119                           f'{{ return "{service.name()}"; }}')
 120
 121         output.write_line()
 122         output.write_line(
 123             '// Used by MethodLookup to identify the generated service base.')
 124         output.write_line(
 125             'constexpr void _PwRpcInternalGeneratedBase() const {}')
 126
 127     service_name_hash = pw_rpc.ids.calculate(service.proto_path())
 128     output.write_line('\n private:')
 129
 130     with output.indent():
 131         output.write_line('friend class ::pw::rpc::internal::MethodLookup;\n')
 132         output.write_line(f'// Hash of "{service.proto_path()}".')
 133         output.write_line(
 134             f'static constexpr uint32_t kServiceId = 0x{service_name_hash:08x};'
 135         )
 136
 137         output.write_line()
 138
 139         # Generate the method table
 140         output.write_line('static constexpr std::array<'
 141                           f'{RPC_NAMESPACE}::internal::{method_union},'
 142                           f' {len(service.methods())}> kMethods = {{')
 143
 144         with output.indent(4):
 145             for method in service.methods():
 146                 method_descriptor(method, pw_rpc.ids.calculate(method.name()),
 147                                   output)
 148
 149         output.write_line('};\n')
 150
 151         # Generate the method lookup table
 152         _method_lookup_table(service, output)
 153
 154     output.write_line('};')
 155
 156     output.write_line('\n}  // namespace generated\n')
 157
 158
 159 def _method_lookup_table(service: ProtoService, output: OutputFile) -> None:
 160     """Generates array of method IDs for looking up methods at compile time."""
 161     output.write_line('static constexpr std::array<uint32_t, '
 162                       f'{len(service.methods())}> kMethodIds = {{')
 163
 164     with output.indent(4):
 165         for method in service.methods():
 166             method_id = pw_rpc.ids.calculate(method.name())
 167             output.write_line(
 168                 f'0x{method_id:08x},  // Hash of "{method.name()}"')
 169
 170     output.write_line('};\n')
 171
 172
 173 StubFunction = Callable[[ProtoServiceMethod, OutputFile], None]
 174
 175 _STUBS_COMMENT = r'''
 176 /*
 177     ____                __                          __        __  _
 178    /  _/___ ___  ____  / /__  ____ ___  ___  ____  / /_____ _/ /_(_)___  ____
 179    / // __ `__ \/ __ \/ / _ \/ __ `__ \/ _ \/ __ \/ __/ __ `/ __/ / __ \/ __ \
 180  _/ // / / / / / /_/ / /  __/ / / / / /  __/ / / / /_/ /_/ / /_/ / /_/ / / / /
 181 /___/_/ /_/ /_/ .___/_/\___/_/ /_/ /_/\___/_/ /_/\__/\__,_/\__/_/\____/_/ /_/
 182              /_/
 183    _____ __        __         __
 184   / ___// /___  __/ /_  _____/ /
 185   \__ \/ __/ / / / __ \/ ___/ /
 186  ___/ / /_/ /_/ / /_/ (__  )_/
 187 /____/\__/\__,_/_.___/____(_)
 188
 189 */
 190 // This section provides stub implementations of the RPC services in this file.
 191 // The code below may be referenced or copied to serve as a starting point for
 192 // your RPC service implementations.
 193 '''
 194
 195
 196 def package_stubs(proto_package: ProtoNode, output: OutputFile,
 197                   unary_stub: StubFunction,
 198                   server_streaming_stub: StubFunction) -> None:
 199
 200     output.write_line('#ifdef _PW_RPC_COMPILE_GENERATED_SERVICE_STUBS')
 201     output.write_line(_STUBS_COMMENT)
 202
 203     output.write_line(f'#include "{output.name()}"\n')
 204
 205     if proto_package.cpp_namespace():
 206         file_namespace = proto_package.cpp_namespace()
 207         if file_namespace.startswith('::'):
 208             file_namespace = file_namespace[2:]
 209
 210         output.write_line(f'namespace {file_namespace} {{')
 211
 212     for node in proto_package:
 213         if node.type() == ProtoNode.Type.SERVICE:
 214             _generate_service_stub(cast(ProtoService, node), output,
 215                                    unary_stub, server_streaming_stub)
 216
 217     if proto_package.cpp_namespace():
 218         output.write_line(f'}}  // namespace {file_namespace}')
 219
 220     output.write_line('\n#endif  // _PW_RPC_COMPILE_GENERATED_SERVICE_STUBS')
 221
 222
 223 def _generate_service_stub(service: ProtoService, output: OutputFile,
 224                            unary_stub: StubFunction,
 225                            server_streaming_stub: StubFunction) -> None:
 226     output.write_line()
 227     output.write_line(
 228         f'class {service.name()} '
 229         f': public generated::{service.name()}<{service.name()}> {{')
 230
 231     output.write_line(' public:')
 232
 233     with output.indent():
 234         blank_line = False
 235
 236         for method in service.methods():
 237             if blank_line:
 238                 output.write_line()
 239             else:
 240                 blank_line = True
 241
 242             if method.type() is ProtoServiceMethod.Type.UNARY:
 243                 unary_stub(method, output)
 244             elif method.type() is ProtoServiceMethod.Type.SERVER_STREAMING:
 245                 server_streaming_stub(method, output)
 246             else:
 247                 raise NotImplementedError(
 248                     'Client and bidirectional streaming not yet implemented')
 249
 250     output.write_line('};\n')