pw_rpc/py/pw_rpc/codegen.py - pigweed/pigweed - Git at Google

 # Copyright 2020 The Pigweed Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not
 # use this file except in compliance with the License. You may obtain a copy of
 # the License at
 #
 #     https://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 # License for the specific language governing permissions and limitations under
 # the License.
 """Common RPC codegen utilities."""

 from datetime import datetime
 import os
 from typing import cast, Any, Callable, Iterable

 from pw_protobuf.output_file import OutputFile
 from pw_protobuf.proto_tree import ProtoNode, ProtoService, ProtoServiceMethod

 import pw_rpc.ids

 PLUGIN_NAME = 'pw_rpc_codegen'
 PLUGIN_VERSION = '0.2.0'

 RPC_NAMESPACE = '::pw::rpc'

 STUB_REQUEST_TODO = (
     '// TODO: Read the request as appropriate for your application')
 STUB_RESPONSE_TODO = (
     '// TODO: Fill in the response as appropriate for your application')
 STUB_WRITER_TODO = (
     '// TODO: Send responses with the writer as appropriate for your '
     'application')

 ServerWriterGenerator = Callable[[OutputFile], None]
 MethodGenerator = Callable[[ProtoServiceMethod, int, OutputFile], None]
 ServiceGenerator = Callable[[ProtoService, ProtoNode, OutputFile], None]
 IncludesGenerator = Callable[[Any, ProtoNode], Iterable[str]]


 def package(file_descriptor_proto, proto_package: ProtoNode,
             output: OutputFile, includes: IncludesGenerator,
             service: ServiceGenerator, client: ServiceGenerator) -> None:
     """Generates service and client code for a package."""
     assert proto_package.type() == ProtoNode.Type.PACKAGE

     output.write_line(f'// {os.path.basename(output.name())} automatically '
                       f'generated by {PLUGIN_NAME} {PLUGIN_VERSION}')
     output.write_line(f'// on {datetime.now().isoformat()}')
     output.write_line('// clang-format off')
     output.write_line('#pragma once\n')

     output.write_line('#include <array>')
     output.write_line('#include <cstdint>')
     output.write_line('#include <type_traits>\n')

     include_lines = [
         '#include "pw_rpc/internal/method_lookup.h"',
         '#include "pw_rpc/server_context.h"',
         '#include "pw_rpc/service.h"',
     ]
     include_lines += includes(file_descriptor_proto, proto_package)

     for include_line in sorted(include_lines):
         output.write_line(include_line)

     output.write_line()

     if proto_package.cpp_namespace():
         file_namespace = proto_package.cpp_namespace()
         if file_namespace.startswith('::'):
             file_namespace = file_namespace[2:]

         output.write_line(f'namespace {file_namespace} {{')

     for node in proto_package:
         if node.type() == ProtoNode.Type.SERVICE:
             service(cast(ProtoService, node), proto_package, output)
             client(cast(ProtoService, node), proto_package, output)

     if proto_package.cpp_namespace():
         output.write_line(f'}}  // namespace {file_namespace}')


 def service_class(service: ProtoService, root: ProtoNode, output: OutputFile,
                   server_writer_alias: ServerWriterGenerator,
                   method_union: str,
                   method_descriptor: MethodGenerator) -> None:
     """Generates a C++ derived class for a nanopb RPC service."""

     output.write_line('namespace generated {')

     base_class = f'{RPC_NAMESPACE}::Service'
     output.write_line('\ntemplate <typename Implementation>')
     output.write_line(
         f'class {service.cpp_namespace(root)} : public {base_class} {{')
     output.write_line(' public:')

     with output.indent():
         output.write_line(
             f'using ServerContext = {RPC_NAMESPACE}::ServerContext;')
         server_writer_alias(output)
         output.write_line()

         output.write_line(f'constexpr {service.name()}()')
         output.write_line(f'    : {base_class}(kServiceId, kMethods) {{}}')

         output.write_line()
         output.write_line(
             f'{service.name()}(const {service.name()}&) = delete;')
         output.write_line(f'{service.name()}& operator='
                           f'(const {service.name()}&) = delete;')

         output.write_line()
         output.write_line(f'static constexpr const char* name() '
                           f'{{ return "{service.name()}"; }}')

         output.write_line()
         output.write_line(
             '// Used by MethodLookup to identify the generated service base.')
         output.write_line(
             'constexpr void _PwRpcInternalGeneratedBase() const {}')

     service_name_hash = pw_rpc.ids.calculate(service.proto_path())
     output.write_line('\n private:')

     with output.indent():
         output.write_line('friend class ::pw::rpc::internal::MethodLookup;\n')
         output.write_line(f'// Hash of "{service.proto_path()}".')
         output.write_line(
             f'static constexpr uint32_t kServiceId = 0x{service_name_hash:08x};'
         )

         output.write_line()

         # Generate the method table
         output.write_line('static constexpr std::array<'
                           f'{RPC_NAMESPACE}::internal::{method_union},'
                           f' {len(service.methods())}> kMethods = {{')

         with output.indent(4):
             for method in service.methods():
                 method_descriptor(method, pw_rpc.ids.calculate(method.name()),
                                   output)

         output.write_line('};\n')

         # Generate the method lookup table
         _method_lookup_table(service, output)

     output.write_line('};')

     output.write_line('\n}  // namespace generated\n')


 def _method_lookup_table(service: ProtoService, output: OutputFile) -> None:
     """Generates array of method IDs for looking up methods at compile time."""
     output.write_line('static constexpr std::array<uint32_t, '
                       f'{len(service.methods())}> kMethodIds = {{')

     with output.indent(4):
         for method in service.methods():
             method_id = pw_rpc.ids.calculate(method.name())
             output.write_line(
                 f'0x{method_id:08x},  // Hash of "{method.name()}"')

     output.write_line('};\n')


 StubFunction = Callable[[ProtoServiceMethod, OutputFile], None]

 _STUBS_COMMENT = r'''
 /*
     ____                __                          __        __  _
    /  _/___ ___  ____  / /__  ____ ___  ___  ____  / /_____ _/ /_(_)___  ____
    / // __ `__ \/ __ \/ / _ \/ __ `__ \/ _ \/ __ \/ __/ __ `/ __/ / __ \/ __ \
  _/ // / / / / / /_/ / /  __/ / / / / /  __/ / / / /_/ /_/ / /_/ / /_/ / / / /
 /___/_/ /_/ /_/ .___/_/\___/_/ /_/ /_/\___/_/ /_/\__/\__,_/\__/_/\____/_/ /_/
              /_/
    _____ __        __         __
   / ___// /___  __/ /_  _____/ /
   \__ \/ __/ / / / __ \/ ___/ /
  ___/ / /_/ /_/ / /_/ (__  )_/
 /____/\__/\__,_/_.___/____(_)

 */
 // This section provides stub implementations of the RPC services in this file.
 // The code below may be referenced or copied to serve as a starting point for
 // your RPC service implementations.
 '''


 def package_stubs(proto_package: ProtoNode, output: OutputFile,
                   unary_stub: StubFunction,
                   server_streaming_stub: StubFunction) -> None:

     output.write_line('#ifdef _PW_RPC_COMPILE_GENERATED_SERVICE_STUBS')
     output.write_line(_STUBS_COMMENT)

     output.write_line(f'#include "{output.name()}"\n')

     if proto_package.cpp_namespace():
         file_namespace = proto_package.cpp_namespace()
         if file_namespace.startswith('::'):
             file_namespace = file_namespace[2:]

         output.write_line(f'namespace {file_namespace} {{')

     for node in proto_package:
         if node.type() == ProtoNode.Type.SERVICE:
             _generate_service_stub(cast(ProtoService, node), output,
                                    unary_stub, server_streaming_stub)

     if proto_package.cpp_namespace():
         output.write_line(f'}}  // namespace {file_namespace}')

     output.write_line('\n#endif  // _PW_RPC_COMPILE_GENERATED_SERVICE_STUBS')


 def _generate_service_stub(service: ProtoService, output: OutputFile,
                            unary_stub: StubFunction,
                            server_streaming_stub: StubFunction) -> None:
     output.write_line()
     output.write_line(
         f'class {service.name()} '
         f': public generated::{service.name()}<{service.name()}> {{')

     output.write_line(' public:')

     with output.indent():
         blank_line = False

         for method in service.methods():
             if blank_line:
                 output.write_line()
             else:
                 blank_line = True

             if method.type() is ProtoServiceMethod.Type.UNARY:
                 unary_stub(method, output)
             elif method.type() is ProtoServiceMethod.Type.SERVER_STREAMING:
                 server_streaming_stub(method, output)
             else:
                 raise NotImplementedError(
                     'Client and bidirectional streaming not yet implemented')

     output.write_line('};\n')
	# Copyright 2020 The Pigweed Authors
	#
	# Licensed under the Apache License, Version 2.0 (the "License"); you may not
	# use this file except in compliance with the License. You may obtain a copy of
	# the License at
	#
	# https://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	# License for the specific language governing permissions and limitations under
	# the License.
	"""Common RPC codegen utilities."""

	from datetime import datetime
	import os
	from typing import cast, Any, Callable, Iterable

	from pw_protobuf.output_file import OutputFile
	from pw_protobuf.proto_tree import ProtoNode, ProtoService, ProtoServiceMethod

	import pw_rpc.ids

	PLUGIN_NAME = 'pw_rpc_codegen'
	PLUGIN_VERSION = '0.2.0'

	RPC_NAMESPACE = '::pw::rpc'

	STUB_REQUEST_TODO = (
	'// TODO: Read the request as appropriate for your application')
	STUB_RESPONSE_TODO = (
	'// TODO: Fill in the response as appropriate for your application')
	STUB_WRITER_TODO = (
	'// TODO: Send responses with the writer as appropriate for your '
	'application')

	ServerWriterGenerator = Callable[[OutputFile], None]
	MethodGenerator = Callable[[ProtoServiceMethod, int, OutputFile], None]
	ServiceGenerator = Callable[[ProtoService, ProtoNode, OutputFile], None]
	IncludesGenerator = Callable[[Any, ProtoNode], Iterable[str]]


	def package(file_descriptor_proto, proto_package: ProtoNode,
	output: OutputFile, includes: IncludesGenerator,
	service: ServiceGenerator, client: ServiceGenerator) -> None:
	"""Generates service and client code for a package."""
	assert proto_package.type() == ProtoNode.Type.PACKAGE

	output.write_line(f'// {os.path.basename(output.name())} automatically '
	f'generated by {PLUGIN_NAME} {PLUGIN_VERSION}')
	output.write_line(f'// on {datetime.now().isoformat()}')
	output.write_line('// clang-format off')
	output.write_line('#pragma once\n')

	output.write_line('#include <array>')
	output.write_line('#include <cstdint>')
	output.write_line('#include <type_traits>\n')

	include_lines = [
	'#include "pw_rpc/internal/method_lookup.h"',
	'#include "pw_rpc/server_context.h"',
	'#include "pw_rpc/service.h"',
	]
	include_lines += includes(file_descriptor_proto, proto_package)

	for include_line in sorted(include_lines):
	output.write_line(include_line)

	output.write_line()

	if proto_package.cpp_namespace():
	file_namespace = proto_package.cpp_namespace()
	if file_namespace.startswith('::'):
	file_namespace = file_namespace[2:]

	output.write_line(f'namespace {file_namespace} {{')

	for node in proto_package:
	if node.type() == ProtoNode.Type.SERVICE:
	service(cast(ProtoService, node), proto_package, output)
	client(cast(ProtoService, node), proto_package, output)

	if proto_package.cpp_namespace():
	output.write_line(f'}} // namespace {file_namespace}')


	def service_class(service: ProtoService, root: ProtoNode, output: OutputFile,
	server_writer_alias: ServerWriterGenerator,
	method_union: str,
	method_descriptor: MethodGenerator) -> None:
	"""Generates a C++ derived class for a nanopb RPC service."""

	output.write_line('namespace generated {')

	base_class = f'{RPC_NAMESPACE}::Service'
	output.write_line('\ntemplate <typename Implementation>')
	output.write_line(
	f'class {service.cpp_namespace(root)} : public {base_class} {{')
	output.write_line(' public:')

	with output.indent():
	output.write_line(
	f'using ServerContext = {RPC_NAMESPACE}::ServerContext;')
	server_writer_alias(output)
	output.write_line()

	output.write_line(f'constexpr {service.name()}()')
	output.write_line(f' : {base_class}(kServiceId, kMethods) {{}}')

	output.write_line()
	output.write_line(
	f'{service.name()}(const {service.name()}&) = delete;')
	output.write_line(f'{service.name()}& operator='
	f'(const {service.name()}&) = delete;')

	output.write_line()
	output.write_line(f'static constexpr const char* name() '
	f'{{ return "{service.name()}"; }}')

	output.write_line()
	output.write_line(
	'// Used by MethodLookup to identify the generated service base.')
	output.write_line(
	'constexpr void _PwRpcInternalGeneratedBase() const {}')

	service_name_hash = pw_rpc.ids.calculate(service.proto_path())
	output.write_line('\n private:')

	with output.indent():
	output.write_line('friend class ::pw::rpc::internal::MethodLookup;\n')
	output.write_line(f'// Hash of "{service.proto_path()}".')
	output.write_line(
	f'static constexpr uint32_t kServiceId = 0x{service_name_hash:08x};'
	)

	output.write_line()

	# Generate the method table
	output.write_line('static constexpr std::array<'
	f'{RPC_NAMESPACE}::internal::{method_union},'
	f' {len(service.methods())}> kMethods = {{')

	with output.indent(4):
	for method in service.methods():
	method_descriptor(method, pw_rpc.ids.calculate(method.name()),
	output)

	output.write_line('};\n')

	# Generate the method lookup table
	_method_lookup_table(service, output)

	output.write_line('};')

	output.write_line('\n} // namespace generated\n')


	def _method_lookup_table(service: ProtoService, output: OutputFile) -> None:
	"""Generates array of method IDs for looking up methods at compile time."""
	output.write_line('static constexpr std::array<uint32_t, '
	f'{len(service.methods())}> kMethodIds = {{')

	with output.indent(4):
	for method in service.methods():
	method_id = pw_rpc.ids.calculate(method.name())
	output.write_line(
	f'0x{method_id:08x}, // Hash of "{method.name()}"')

	output.write_line('};\n')


	StubFunction = Callable[[ProtoServiceMethod, OutputFile], None]

	_STUBS_COMMENT = r'''
	/*
	____ __ __ __ _
	/ _/___ ___ ____ / /__ ____ ___ ___ ____ / /_____ _/ /_(_)___ ____
	/ // __ `__ \/ __ \/ / _ \/ __ `__ \/ _ \/ __ \/ __/ __ `/ __/ / __ \/ __ \
	_/ // / / / / / /_/ / / __/ / / / / / __/ / / / /_/ /_/ / /_/ / /_/ / / / /
	/___/_/ /_/ /_/ .___/_/\___/_/ /_/ /_/\___/_/ /_/\__/\__,_/\__/_/\____/_/ /_/
	/_/
	_____ __ __ __
	/ ___// /___ __/ /_ _____/ /
	\__ \/ __/ / / / __ \/ ___/ /
	___/ / /_/ /_/ / /_/ (__ )_/
	/____/\__/\__,_/_.___/____(_)

	*/
	// This section provides stub implementations of the RPC services in this file.
	// The code below may be referenced or copied to serve as a starting point for
	// your RPC service implementations.
	'''


	def package_stubs(proto_package: ProtoNode, output: OutputFile,
	unary_stub: StubFunction,
	server_streaming_stub: StubFunction) -> None:

	output.write_line('#ifdef _PW_RPC_COMPILE_GENERATED_SERVICE_STUBS')
	output.write_line(_STUBS_COMMENT)

	output.write_line(f'#include "{output.name()}"\n')

	if proto_package.cpp_namespace():
	file_namespace = proto_package.cpp_namespace()
	if file_namespace.startswith('::'):
	file_namespace = file_namespace[2:]

	output.write_line(f'namespace {file_namespace} {{')

	for node in proto_package:
	if node.type() == ProtoNode.Type.SERVICE:
	_generate_service_stub(cast(ProtoService, node), output,
	unary_stub, server_streaming_stub)

	if proto_package.cpp_namespace():
	output.write_line(f'}} // namespace {file_namespace}')

	output.write_line('\n#endif // _PW_RPC_COMPILE_GENERATED_SERVICE_STUBS')


	def _generate_service_stub(service: ProtoService, output: OutputFile,
	unary_stub: StubFunction,
	server_streaming_stub: StubFunction) -> None:
	output.write_line()
	output.write_line(
	f'class {service.name()} '
	f': public generated::{service.name()}<{service.name()}> {{')

	output.write_line(' public:')

	with output.indent():
	blank_line = False

	for method in service.methods():
	if blank_line:
	output.write_line()
	else:
	blank_line = True

	if method.type() is ProtoServiceMethod.Type.UNARY:
	unary_stub(method, output)
	elif method.type() is ProtoServiceMethod.Type.SERVER_STREAMING:
	server_streaming_stub(method, output)
	else:
	raise NotImplementedError(
	'Client and bidirectional streaming not yet implemented')

	output.write_line('};\n')