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

 #!/usr/bin/env python3
 # Copyright 2019 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.
 """pw_protobuf compiler plugin.

 This file implements a protobuf compiler plugin which generates C++ headers for
 protobuf messages in the pw_protobuf format.
 """

 import os
 import sys

 from typing import List

 import google.protobuf.compiler.plugin_pb2 as plugin_pb2
 import google.protobuf.descriptor_pb2 as descriptor_pb2

 from pw_protobuf.methods import PROTO_FIELD_METHODS
 from pw_protobuf.proto_structures import ProtoEnum, ProtoMessage
 from pw_protobuf.proto_structures import ProtoMessageField, ProtoNode
 from pw_protobuf.proto_structures import ProtoPackage

 PLUGIN_NAME = 'pw_protobuf'
 PLUGIN_VERSION = '0.1.0'

 PROTOBUF_NAMESPACE = 'pw::protobuf'
 BASE_PROTO_CLASS = 'ProtoMessageEncoder'


 # protoc captures stdout, so we need to printf debug to stderr.
 def debug_print(*args, **kwargs):
     print(*args, file=sys.stderr, **kwargs)


 class OutputFile:
     """A buffer to which data is written.

     Example:

     ```
     output = Output("hello.c")
     output.write_line('int main(void) {')
     with output.indent():
         output.write_line('printf("Hello, world");')
         output.write_line('return 0;')
     output.write_line('}')

     print(output.content())
     ```

     Produces:
     ```
     int main(void) {
       printf("Hello, world");
       return 0;
     }
     ```
     """

     INDENT_WIDTH = 2

     def __init__(self, filename: str):
         self._filename: str = filename
         self._content: List[str] = []
         self._indentation: int = 0

     def write_line(self, line: str = '') -> None:
         if line:
             self._content.append(' ' * self._indentation)
             self._content.append(line)
         self._content.append('\n')

     def indent(self) -> 'OutputFile._IndentationContext':
         """Increases the indentation level of the output."""
         return self._IndentationContext(self)

     def name(self) -> str:
         return self._filename

     def content(self) -> str:
         return ''.join(self._content)

     class _IndentationContext:
         """Context that increases the output's indentation when it is active."""
         def __init__(self, output: 'OutputFile'):
             self._output = output

         def __enter__(self):
             self._output._indentation += OutputFile.INDENT_WIDTH

         def __exit__(self, typ, value, traceback):
             self._output._indentation -= OutputFile.INDENT_WIDTH


 def generate_code_for_message(
         message: ProtoNode,
         root: ProtoNode,
         output: OutputFile,
 ) -> None:
     """Creates a C++ class for a protobuf message."""
     assert message.type() == ProtoNode.Type.MESSAGE

     # Message classes inherit from the base proto message class in codegen.h
     # and use its constructor.
     base_class = f'{PROTOBUF_NAMESPACE}::{BASE_PROTO_CLASS}'
     output.write_line(
         f'class {message.cpp_namespace(root)}::Encoder : public {base_class} {{'
     )
     output.write_line(' public:')

     with output.indent():
         output.write_line(f'using {BASE_PROTO_CLASS}::{BASE_PROTO_CLASS};')

         # Generate methods for each of the message's fields.
         for field in message.fields():
             for method_class in PROTO_FIELD_METHODS[field.type()]:
                 method = method_class(field, message, root)
                 if not method.should_appear():
                     continue

                 output.write_line()
                 method_signature = (
                     f'{method.return_type()} '
                     f'{method.name()}({method.param_string()})')

                 if not method.in_class_definition():
                     # Method will be defined outside of the class at the end of
                     # the file.
                     output.write_line(f'{method_signature};')
                     continue

                 output.write_line(f'{method_signature} {{')
                 with output.indent():
                     for line in method.body():
                         output.write_line(line)
                 output.write_line('}')

     output.write_line('};')


 def define_not_in_class_methods(
         message: ProtoNode,
         root: ProtoNode,
         output: OutputFile,
 ) -> None:
     """Defines methods for a message class that were previously declared."""
     assert message.type() == ProtoNode.Type.MESSAGE

     for field in message.fields():
         for method_class in PROTO_FIELD_METHODS[field.type()]:
             method = method_class(field, message, root)
             if not method.should_appear() or method.in_class_definition():
                 continue

             output.write_line()
             class_name = f'{message.cpp_namespace(root)}::Encoder'
             method_signature = (
                 f'inline {method.return_type(from_root=True)} '
                 f'{class_name}::{method.name()}({method.param_string()})')
             output.write_line(f'{method_signature} {{')
             with output.indent():
                 for line in method.body():
                     output.write_line(line)
             output.write_line('}')


 def generate_code_for_enum(
         enum: ProtoNode,
         root: ProtoNode,
         output: OutputFile,
 ) -> None:
     """Creates a C++ enum for a proto enum."""
     assert enum.type() == ProtoNode.Type.ENUM

     output.write_line(f'enum class {enum.cpp_namespace(root)} {{')
     with output.indent():
         for name, number in enum.values():
             output.write_line(f'{name} = {number},')
     output.write_line('};')


 def forward_declare(
         node: ProtoNode,
         root: ProtoNode,
         output: OutputFile,
 ) -> None:
     """Generates code forward-declaring entities in a message's namespace."""
     if node.type() != ProtoNode.Type.MESSAGE:
         return

     namespace = node.cpp_namespace(root)
     output.write_line()
     output.write_line(f'namespace {namespace} {{')

     # Define an enum defining each of the message's fields and their numbers.
     output.write_line('enum class Fields {')
     with output.indent():
         for field in node.fields():
             output.write_line(f'{field.enum_name()} = {field.number()},')
     output.write_line('};')

     # Declare the message's encoder class and all of its enums.
     output.write_line()
     output.write_line('class Encoder;')
     for child in node.children():
         if child.type() == ProtoNode.Type.ENUM:
             output.write_line()
             generate_code_for_enum(child, node, output)

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


 # TODO(frolv): Right now, this plugin assumes that package is synonymous with
 # .proto file. This will need to be updated to handle compiling multiple files.
 def generate_code_for_package(package: ProtoNode, output: OutputFile) -> None:
     """Generates code for a single .pb.h file corresponding to a .proto file."""

     assert 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('#pragma once\n')
     output.write_line('#include <cstddef>')
     output.write_line('#include <cstdint>\n')
     output.write_line('#include "pw_protobuf/codegen.h"')

     if package.cpp_namespace():
         output.write_line(f'\nnamespace {package.cpp_namespace()} {{')

     for node in package:
         forward_declare(node, package, output)

     # Define all top-level enums.
     for node in package.children():
         if node.type() == ProtoNode.Type.ENUM:
             output.write_line()
             generate_code_for_enum(node, package, output)

     # Run through all messages in the file, generating a class for each.
     for node in package:
         if node.type() == ProtoNode.Type.MESSAGE:
             output.write_line()
             generate_code_for_message(node, package, output)

     # Run a second pass through the classes, this time defining all of the
     # methods which were previously only declared.
     for node in package:
         if node.type() == ProtoNode.Type.MESSAGE:
             define_not_in_class_methods(node, package, output)

     if package.cpp_namespace():
         output.write_line(f'\n}}  // namespace {package.cpp_namespace()}')


 def add_enum_fields(enum: ProtoNode, proto_enum) -> None:
     """Adds fields from a protobuf enum descriptor to an enum node."""
     assert enum.type() == ProtoNode.Type.ENUM
     for value in proto_enum.value:
         enum.add_value(value.name, value.number)


 def add_message_fields(
         root: ProtoNode,
         message: ProtoNode,
         proto_message,
 ) -> None:
     """Adds fields from a protobuf message descriptor to a message node."""
     assert message.type() == ProtoNode.Type.MESSAGE

     for field in proto_message.field:
         if field.type_name:
             # The "type_name" member contains the global .proto path of the
             # field's type object, for example ".pw.protobuf.test.KeyValuePair".
             # Since only a single proto file is currently supported, the root
             # node has the value of the file's package ("pw.protobuf.test").
             # This must be stripped from the path to find the desired node
             # within the tree.
             #
             # TODO(frolv): Once multiple files are supported, the root node
             # should refer to the global namespace, and this should no longer
             # be needed.
             path = field.type_name
             if path[0] == '.':
                 path = path[1:]

             if path.startswith(root.name()):
                 relative_path = path[len(root.name()):].lstrip('.')
             else:
                 relative_path = path

             type_node = root.find(relative_path)
         else:
             type_node = None

         repeated = \
             field.label == descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED
         message.add_field(
             ProtoMessageField(
                 field.name,
                 field.number,
                 field.type,
                 type_node,
                 repeated,
             ))


 def populate_fields(proto_file, root: ProtoNode) -> None:
     """Traverses a proto file, adding all message and enum fields to a tree."""
     def populate_message(node, message):
         """Recursively populates nested messages and enums."""
         add_message_fields(root, node, message)

         for enum in message.enum_type:
             add_enum_fields(node.find(enum.name), enum)
         for msg in message.nested_type:
             populate_message(node.find(msg.name), msg)

     # Iterate through the proto file, populating top-level enums and messages.
     for enum in proto_file.enum_type:
         add_enum_fields(root.find(enum.name), enum)
     for message in proto_file.message_type:
         populate_message(root.find(message.name), message)


 def build_hierarchy(proto_file):
     """Creates a ProtoNode hierarchy from a proto file descriptor."""

     root = ProtoPackage(proto_file.package)

     def build_message_subtree(proto_message):
         node = ProtoMessage(proto_message.name)
         for enum in proto_message.enum_type:
             node.add_child(ProtoEnum(enum.name))
         for submessage in proto_message.nested_type:
             node.add_child(build_message_subtree(submessage))

         return node

     for enum in proto_file.enum_type:
         root.add_child(ProtoEnum(enum.name))

     for message in proto_file.message_type:
         root.add_child(build_message_subtree(message))

     return root


 def process_proto_file(proto_file):
     """Generates code for a single .proto file."""

     # Two passes are made through the file. The first builds the tree of all
     # message/enum nodes, then the second creates the fields in each. This is
     # done as non-primitive fields need pointers to their types, which requires
     # the entire tree to have been parsed into memory.
     root = build_hierarchy(proto_file)
     populate_fields(proto_file, root)

     output_filename = os.path.splitext(proto_file.name)[0] + '.pb.h'
     output_file = OutputFile(output_filename)
     generate_code_for_package(root, output_file)

     return output_file


 def process_proto_request(req: plugin_pb2.CodeGeneratorRequest,
                           res: plugin_pb2.CodeGeneratorResponse) -> None:
     """Handles a protoc CodeGeneratorRequest message.

     Generates code for the files in the request and writes the output to the
     specified CodeGeneratorResponse message.

     Args:
       req: A CodeGeneratorRequest for a proto compilation.
       res: A CodeGeneratorResponse to populate with the plugin's output.
     """
     for proto_file in req.proto_file:
         # TODO(frolv): Proto files are currently processed individually. Support
         # for multiple files with cross-dependencies should be added.
         output_file = process_proto_file(proto_file)
         fd = res.file.add()
         fd.name = output_file.name()
         fd.content = output_file.content()


 def main() -> int:
     """Protobuf compiler plugin entrypoint.

     Reads a CodeGeneratorRequest proto from stdin and writes a
     CodeGeneratorResponse to stdout.
     """
     data = sys.stdin.buffer.read()
     request = plugin_pb2.CodeGeneratorRequest.FromString(data)
     response = plugin_pb2.CodeGeneratorResponse()
     process_proto_request(request, response)
     sys.stdout.buffer.write(response.SerializeToString())
     return 0


 if __name__ == '__main__':
     sys.exit(main())
	#!/usr/bin/env python3
	# Copyright 2019 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.
	"""pw_protobuf compiler plugin.

	This file implements a protobuf compiler plugin which generates C++ headers for
	protobuf messages in the pw_protobuf format.
	"""

	import os
	import sys

	from typing import List

	import google.protobuf.compiler.plugin_pb2 as plugin_pb2
	import google.protobuf.descriptor_pb2 as descriptor_pb2

	from pw_protobuf.methods import PROTO_FIELD_METHODS
	from pw_protobuf.proto_structures import ProtoEnum, ProtoMessage
	from pw_protobuf.proto_structures import ProtoMessageField, ProtoNode
	from pw_protobuf.proto_structures import ProtoPackage

	PLUGIN_NAME = 'pw_protobuf'
	PLUGIN_VERSION = '0.1.0'

	PROTOBUF_NAMESPACE = 'pw::protobuf'
	BASE_PROTO_CLASS = 'ProtoMessageEncoder'


	# protoc captures stdout, so we need to printf debug to stderr.
	def debug_print(args, *kwargs):
	print(args, file=sys.stderr, *kwargs)


	class OutputFile:
	"""A buffer to which data is written.

	Example:

	```
	output = Output("hello.c")
	output.write_line('int main(void) {')
	with output.indent():
	output.write_line('printf("Hello, world");')
	output.write_line('return 0;')
	output.write_line('}')

	print(output.content())
	```

	Produces:
	```
	int main(void) {
	printf("Hello, world");
	return 0;
	}
	```
	"""

	INDENT_WIDTH = 2

	def __init__(self, filename: str):
	self._filename: str = filename
	self._content: List[str] = []
	self._indentation: int = 0

	def write_line(self, line: str = '') -> None:
	if line:
	self._content.append(' ' * self._indentation)
	self._content.append(line)
	self._content.append('\n')

	def indent(self) -> 'OutputFile._IndentationContext':
	"""Increases the indentation level of the output."""
	return self._IndentationContext(self)

	def name(self) -> str:
	return self._filename

	def content(self) -> str:
	return ''.join(self._content)

	class _IndentationContext:
	"""Context that increases the output's indentation when it is active."""
	def __init__(self, output: 'OutputFile'):
	self._output = output

	def __enter__(self):
	self._output._indentation += OutputFile.INDENT_WIDTH

	def __exit__(self, typ, value, traceback):
	self._output._indentation -= OutputFile.INDENT_WIDTH


	def generate_code_for_message(
	message: ProtoNode,
	root: ProtoNode,
	output: OutputFile,
	) -> None:
	"""Creates a C++ class for a protobuf message."""
	assert message.type() == ProtoNode.Type.MESSAGE

	# Message classes inherit from the base proto message class in codegen.h
	# and use its constructor.
	base_class = f'{PROTOBUF_NAMESPACE}::{BASE_PROTO_CLASS}'
	output.write_line(
	f'class {message.cpp_namespace(root)}::Encoder : public {base_class} {{'
	)
	output.write_line(' public:')

	with output.indent():
	output.write_line(f'using {BASE_PROTO_CLASS}::{BASE_PROTO_CLASS};')

	# Generate methods for each of the message's fields.
	for field in message.fields():
	for method_class in PROTO_FIELD_METHODS[field.type()]:
	method = method_class(field, message, root)
	if not method.should_appear():
	continue

	output.write_line()
	method_signature = (
	f'{method.return_type()} '
	f'{method.name()}({method.param_string()})')

	if not method.in_class_definition():
	# Method will be defined outside of the class at the end of
	# the file.
	output.write_line(f'{method_signature};')
	continue

	output.write_line(f'{method_signature} {{')
	with output.indent():
	for line in method.body():
	output.write_line(line)
	output.write_line('}')

	output.write_line('};')


	def define_not_in_class_methods(
	message: ProtoNode,
	root: ProtoNode,
	output: OutputFile,
	) -> None:
	"""Defines methods for a message class that were previously declared."""
	assert message.type() == ProtoNode.Type.MESSAGE

	for field in message.fields():
	for method_class in PROTO_FIELD_METHODS[field.type()]:
	method = method_class(field, message, root)
	if not method.should_appear() or method.in_class_definition():
	continue

	output.write_line()
	class_name = f'{message.cpp_namespace(root)}::Encoder'
	method_signature = (
	f'inline {method.return_type(from_root=True)} '
	f'{class_name}::{method.name()}({method.param_string()})')
	output.write_line(f'{method_signature} {{')
	with output.indent():
	for line in method.body():
	output.write_line(line)
	output.write_line('}')


	def generate_code_for_enum(
	enum: ProtoNode,
	root: ProtoNode,
	output: OutputFile,
	) -> None:
	"""Creates a C++ enum for a proto enum."""
	assert enum.type() == ProtoNode.Type.ENUM

	output.write_line(f'enum class {enum.cpp_namespace(root)} {{')
	with output.indent():
	for name, number in enum.values():
	output.write_line(f'{name} = {number},')
	output.write_line('};')


	def forward_declare(
	node: ProtoNode,
	root: ProtoNode,
	output: OutputFile,
	) -> None:
	"""Generates code forward-declaring entities in a message's namespace."""
	if node.type() != ProtoNode.Type.MESSAGE:
	return

	namespace = node.cpp_namespace(root)
	output.write_line()
	output.write_line(f'namespace {namespace} {{')

	# Define an enum defining each of the message's fields and their numbers.
	output.write_line('enum class Fields {')
	with output.indent():
	for field in node.fields():
	output.write_line(f'{field.enum_name()} = {field.number()},')
	output.write_line('};')

	# Declare the message's encoder class and all of its enums.
	output.write_line()
	output.write_line('class Encoder;')
	for child in node.children():
	if child.type() == ProtoNode.Type.ENUM:
	output.write_line()
	generate_code_for_enum(child, node, output)

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


	# TODO(frolv): Right now, this plugin assumes that package is synonymous with
	# .proto file. This will need to be updated to handle compiling multiple files.
	def generate_code_for_package(package: ProtoNode, output: OutputFile) -> None:
	"""Generates code for a single .pb.h file corresponding to a .proto file."""

	assert 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('#pragma once\n')
	output.write_line('#include <cstddef>')
	output.write_line('#include <cstdint>\n')
	output.write_line('#include "pw_protobuf/codegen.h"')

	if package.cpp_namespace():
	output.write_line(f'\nnamespace {package.cpp_namespace()} {{')

	for node in package:
	forward_declare(node, package, output)

	# Define all top-level enums.
	for node in package.children():
	if node.type() == ProtoNode.Type.ENUM:
	output.write_line()
	generate_code_for_enum(node, package, output)

	# Run through all messages in the file, generating a class for each.
	for node in package:
	if node.type() == ProtoNode.Type.MESSAGE:
	output.write_line()
	generate_code_for_message(node, package, output)

	# Run a second pass through the classes, this time defining all of the
	# methods which were previously only declared.
	for node in package:
	if node.type() == ProtoNode.Type.MESSAGE:
	define_not_in_class_methods(node, package, output)

	if package.cpp_namespace():
	output.write_line(f'\n}} // namespace {package.cpp_namespace()}')


	def add_enum_fields(enum: ProtoNode, proto_enum) -> None:
	"""Adds fields from a protobuf enum descriptor to an enum node."""
	assert enum.type() == ProtoNode.Type.ENUM
	for value in proto_enum.value:
	enum.add_value(value.name, value.number)


	def add_message_fields(
	root: ProtoNode,
	message: ProtoNode,
	proto_message,
	) -> None:
	"""Adds fields from a protobuf message descriptor to a message node."""
	assert message.type() == ProtoNode.Type.MESSAGE

	for field in proto_message.field:
	if field.type_name:
	# The "type_name" member contains the global .proto path of the
	# field's type object, for example ".pw.protobuf.test.KeyValuePair".
	# Since only a single proto file is currently supported, the root
	# node has the value of the file's package ("pw.protobuf.test").
	# This must be stripped from the path to find the desired node
	# within the tree.
	#
	# TODO(frolv): Once multiple files are supported, the root node
	# should refer to the global namespace, and this should no longer
	# be needed.
	path = field.type_name
	if path[0] == '.':
	path = path[1:]

	if path.startswith(root.name()):
	relative_path = path[len(root.name()):].lstrip('.')
	else:
	relative_path = path

	type_node = root.find(relative_path)
	else:
	type_node = None

	repeated = \
	field.label == descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED
	message.add_field(
	ProtoMessageField(
	field.name,
	field.number,
	field.type,
	type_node,
	repeated,
	))


	def populate_fields(proto_file, root: ProtoNode) -> None:
	"""Traverses a proto file, adding all message and enum fields to a tree."""
	def populate_message(node, message):
	"""Recursively populates nested messages and enums."""
	add_message_fields(root, node, message)

	for enum in message.enum_type:
	add_enum_fields(node.find(enum.name), enum)
	for msg in message.nested_type:
	populate_message(node.find(msg.name), msg)

	# Iterate through the proto file, populating top-level enums and messages.
	for enum in proto_file.enum_type:
	add_enum_fields(root.find(enum.name), enum)
	for message in proto_file.message_type:
	populate_message(root.find(message.name), message)


	def build_hierarchy(proto_file):
	"""Creates a ProtoNode hierarchy from a proto file descriptor."""

	root = ProtoPackage(proto_file.package)

	def build_message_subtree(proto_message):
	node = ProtoMessage(proto_message.name)
	for enum in proto_message.enum_type:
	node.add_child(ProtoEnum(enum.name))
	for submessage in proto_message.nested_type:
	node.add_child(build_message_subtree(submessage))

	return node

	for enum in proto_file.enum_type:
	root.add_child(ProtoEnum(enum.name))

	for message in proto_file.message_type:
	root.add_child(build_message_subtree(message))

	return root


	def process_proto_file(proto_file):
	"""Generates code for a single .proto file."""

	# Two passes are made through the file. The first builds the tree of all
	# message/enum nodes, then the second creates the fields in each. This is
	# done as non-primitive fields need pointers to their types, which requires
	# the entire tree to have been parsed into memory.
	root = build_hierarchy(proto_file)
	populate_fields(proto_file, root)

	output_filename = os.path.splitext(proto_file.name)[0] + '.pb.h'
	output_file = OutputFile(output_filename)
	generate_code_for_package(root, output_file)

	return output_file


	def process_proto_request(req: plugin_pb2.CodeGeneratorRequest,
	res: plugin_pb2.CodeGeneratorResponse) -> None:
	"""Handles a protoc CodeGeneratorRequest message.

	Generates code for the files in the request and writes the output to the
	specified CodeGeneratorResponse message.

	Args:
	req: A CodeGeneratorRequest for a proto compilation.
	res: A CodeGeneratorResponse to populate with the plugin's output.
	"""
	for proto_file in req.proto_file:
	# TODO(frolv): Proto files are currently processed individually. Support
	# for multiple files with cross-dependencies should be added.
	output_file = process_proto_file(proto_file)
	fd = res.file.add()
	fd.name = output_file.name()
	fd.content = output_file.content()


	def main() -> int:
	"""Protobuf compiler plugin entrypoint.

	Reads a CodeGeneratorRequest proto from stdin and writes a
	CodeGeneratorResponse to stdout.
	"""
	data = sys.stdin.buffer.read()
	request = plugin_pb2.CodeGeneratorRequest.FromString(data)
	response = plugin_pb2.CodeGeneratorResponse()
	process_proto_request(request, response)
	sys.stdout.buffer.write(response.SerializeToString())
	return 0


	if __name__ == '__main__':
	sys.exit(main())