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

 # 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.
 """This module defines data structures for protobuf entities."""

 import abc
 import collections
 import enum

 from typing import Callable, Dict, Iterator, List, Optional, Tuple, TypeVar

 T = TypeVar('T')  # pylint: disable=invalid-name


 class ProtoNode(abc.ABC):
     """A ProtoNode represents a C++ scope mapping of an entity in a .proto file.

     Nodes form a tree beginning at a top-level (global) scope, descending into a
     hierarchy of .proto packages and the messages and enums defined within them.
     """
     class Type(enum.Enum):
         """The type of a ProtoNode.

         PACKAGE maps to a C++ namespace.
         MESSAGE maps to a C++ "Encoder" class within its own namespace.
         ENUM maps to a C++ enum within its parent's namespace.
         """
         PACKAGE = 1
         MESSAGE = 2
         ENUM = 3

     def __init__(self, name: str):
         self._name: str = name
         self._children: Dict[str, 'ProtoNode'] = collections.OrderedDict()
         self._parent: Optional['ProtoNode'] = None

     @abc.abstractmethod
     def type(self) -> 'ProtoNode.Type':
         """The type of the node."""

     def children(self) -> List['ProtoNode']:
         return self._children.values()

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

     def cpp_name(self) -> str:
         """The name of this node in generated C++ code."""
         return self._name.replace('.', '::')

     def cpp_namespace(self, root: Optional['ProtoNode'] = None) -> str:
         """C++ namespace of the node, up to the specified root."""
         return '::'.join(
             self._attr_hierarchy(lambda node: node.cpp_name(), root))

     def common_ancestor(self, other: 'ProtoNode') -> Optional['ProtoNode']:
         """Finds the earliest common ancestor of this node and other."""

         if other is None:
             return None

         own_depth = self.depth()
         other_depth = other.depth()
         diff = abs(own_depth - other_depth)

         if own_depth < other_depth:
             first: Optional['ProtoNode'] = self
             second: Optional['ProtoNode'] = other
         else:
             first = other
             second = self

         while diff > 0:
             second = second.parent()
             diff -= 1

         while first != second:
             if first is None or second is None:
                 return None

             first = first.parent()
             second = second.parent()

         return first

     def depth(self) -> int:
         """Returns the depth of this node from the root."""
         depth = 0
         node = self._parent
         while node:
             depth += 1
             node = node.parent()
         return depth

     def add_child(self, child: 'ProtoNode') -> None:
         """Inserts a new node into the tree as a child of this node.

         Args:
           child: The node to insert.

         Raises:
           ValueError: This node does not allow nesting the given type of child.
         """
         if not self._supports_child(child):
             raise ValueError('Invalid child %s for node of type %s' %
                              (child.type(), self.type()))

         # pylint: disable=protected-access
         if child.parent() is not None:
             del child._parent._children[child.name()]

         child._parent = self
         self._children[child.name()] = child
         # pylint: enable=protected-access

     def find(self, path: str) -> Optional['ProtoNode']:
         """Finds a node within this node's subtree."""
         node = self

         # pylint: disable=protected-access
         for section in path.split('.'):
             node = node._children.get(section)
             if node is None:
                 return None
         # pylint: enable=protected-access

         return node

     def parent(self) -> Optional['ProtoNode']:
         return self._parent

     def __iter__(self) -> Iterator['ProtoNode']:
         """Iterates depth-first through all nodes in this node's subtree."""
         yield self
         for child_iterator in self._children.values():
             for child in child_iterator:
                 yield child

     def _attr_hierarchy(
             self,
             attr_accessor: Callable[['ProtoNode'], T],
             root: Optional['ProtoNode'],
     ) -> Iterator[T]:
         """Fetches node attributes at each level of the tree from the root.

         Args:
           attr_accessor: Function which extracts attributes from a ProtoNode.
           root: The node at which to terminate.

         Returns:
           An iterator to a list of the selected attributes from the root to the
           current node.
         """
         hierarchy = []
         node: Optional['ProtoNode'] = self
         while node is not None and node != root:
             hierarchy.append(attr_accessor(node))
             node = node.parent()
         return reversed(hierarchy)

     @abc.abstractmethod
     def _supports_child(self, child: 'ProtoNode') -> bool:
         """Returns True if child is a valid child type for the current node."""


 class ProtoPackage(ProtoNode):
     """A protobuf package."""
     def type(self) -> ProtoNode.Type:
         return ProtoNode.Type.PACKAGE

     def _supports_child(self, child: ProtoNode) -> bool:
         return True


 class ProtoEnum(ProtoNode):
     """Representation of an enum in a .proto file."""

     # Prefix for names of values in C++ enums.
     ENUM_PREFIX: str = 'k'

     def __init__(self, name: str):
         super().__init__(name)
         self._values: List[Tuple[str, int]] = []

     def type(self) -> ProtoNode.Type:
         return ProtoNode.Type.ENUM

     def values(self) -> List[Tuple[str, int]]:
         return list(self._values)

     def add_value(self, name: str, value: int) -> None:
         name = '{}{}'.format(ProtoEnum.ENUM_PREFIX,
                              ProtoMessageField.canonicalize_name(name))
         self._values.append((name, value))

     def _supports_child(self, child: ProtoNode) -> bool:
         # Enums cannot have nested children.
         return False


 class ProtoMessage(ProtoNode):
     """Representation of a message in a .proto file."""
     def __init__(self, name: str):
         super().__init__(name)
         self._fields: List['ProtoMessageField'] = []

     def type(self) -> ProtoNode.Type:
         return ProtoNode.Type.MESSAGE

     def fields(self) -> List['ProtoMessageField']:
         return list(self._fields)

     def add_field(self, field: 'ProtoMessageField') -> None:
         self._fields.append(field)

     def _supports_child(self, child: ProtoNode) -> bool:
         return (child.type() == self.Type.ENUM
                 or child.type() == self.Type.MESSAGE)


 # This class is not a node and does not appear in the proto tree.
 # Fields belong to proto messages and are processed separately.
 class ProtoMessageField:
     """Representation of a field within a protobuf message."""
     def __init__(self,
                  field_name: str,
                  field_number: int,
                  field_type: int,
                  type_node: Optional[ProtoNode] = None,
                  repeated: bool = False):
         self._name: str = self.canonicalize_name(field_name)
         self._number: int = field_number
         self._type: int = field_type
         self._type_node: Optional[ProtoNode] = type_node
         self._repeated: bool = repeated

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

     def enum_name(self) -> str:
         return '{}{}'.format(ProtoEnum.ENUM_PREFIX, self._name)

     def number(self) -> int:
         return self._number

     def type(self) -> int:
         return self._type

     def type_node(self) -> Optional[ProtoNode]:
         return self._type_node

     def is_repeated(self) -> bool:
         return self._repeated

     @staticmethod
     def canonicalize_name(field_name: str) -> str:
         """Converts a field name to UpperCamelCase."""
         name_components = field_name.split('_')
         for i, _ in enumerate(name_components):
             name_components[i] = name_components[i].lower().capitalize()
         return ''.join(name_components)
	# 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.
	"""This module defines data structures for protobuf entities."""

	import abc
	import collections
	import enum

	from typing import Callable, Dict, Iterator, List, Optional, Tuple, TypeVar

	T = TypeVar('T') # pylint: disable=invalid-name


	class ProtoNode(abc.ABC):
	"""A ProtoNode represents a C++ scope mapping of an entity in a .proto file.

	Nodes form a tree beginning at a top-level (global) scope, descending into a
	hierarchy of .proto packages and the messages and enums defined within them.
	"""
	class Type(enum.Enum):
	"""The type of a ProtoNode.

	PACKAGE maps to a C++ namespace.
	MESSAGE maps to a C++ "Encoder" class within its own namespace.
	ENUM maps to a C++ enum within its parent's namespace.
	"""
	PACKAGE = 1
	MESSAGE = 2
	ENUM = 3

	def __init__(self, name: str):
	self._name: str = name
	self._children: Dict[str, 'ProtoNode'] = collections.OrderedDict()
	self._parent: Optional['ProtoNode'] = None

	@abc.abstractmethod
	def type(self) -> 'ProtoNode.Type':
	"""The type of the node."""

	def children(self) -> List['ProtoNode']:
	return self._children.values()

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

	def cpp_name(self) -> str:
	"""The name of this node in generated C++ code."""
	return self._name.replace('.', '::')

	def cpp_namespace(self, root: Optional['ProtoNode'] = None) -> str:
	"""C++ namespace of the node, up to the specified root."""
	return '::'.join(
	self._attr_hierarchy(lambda node: node.cpp_name(), root))

	def common_ancestor(self, other: 'ProtoNode') -> Optional['ProtoNode']:
	"""Finds the earliest common ancestor of this node and other."""

	if other is None:
	return None

	own_depth = self.depth()
	other_depth = other.depth()
	diff = abs(own_depth - other_depth)

	if own_depth < other_depth:
	first: Optional['ProtoNode'] = self
	second: Optional['ProtoNode'] = other
	else:
	first = other
	second = self

	while diff > 0:
	second = second.parent()
	diff -= 1

	while first != second:
	if first is None or second is None:
	return None

	first = first.parent()
	second = second.parent()

	return first

	def depth(self) -> int:
	"""Returns the depth of this node from the root."""
	depth = 0
	node = self._parent
	while node:
	depth += 1
	node = node.parent()
	return depth

	def add_child(self, child: 'ProtoNode') -> None:
	"""Inserts a new node into the tree as a child of this node.

	Args:
	child: The node to insert.

	Raises:
	ValueError: This node does not allow nesting the given type of child.
	"""
	if not self._supports_child(child):
	raise ValueError('Invalid child %s for node of type %s' %
	(child.type(), self.type()))

	# pylint: disable=protected-access
	if child.parent() is not None:
	del child._parent._children[child.name()]

	child._parent = self
	self._children[child.name()] = child
	# pylint: enable=protected-access

	def find(self, path: str) -> Optional['ProtoNode']:
	"""Finds a node within this node's subtree."""
	node = self

	# pylint: disable=protected-access
	for section in path.split('.'):
	node = node._children.get(section)
	if node is None:
	return None
	# pylint: enable=protected-access

	return node

	def parent(self) -> Optional['ProtoNode']:
	return self._parent

	def __iter__(self) -> Iterator['ProtoNode']:
	"""Iterates depth-first through all nodes in this node's subtree."""
	yield self
	for child_iterator in self._children.values():
	for child in child_iterator:
	yield child

	def _attr_hierarchy(
	self,
	attr_accessor: Callable[['ProtoNode'], T],
	root: Optional['ProtoNode'],
	) -> Iterator[T]:
	"""Fetches node attributes at each level of the tree from the root.

	Args:
	attr_accessor: Function which extracts attributes from a ProtoNode.
	root: The node at which to terminate.

	Returns:
	An iterator to a list of the selected attributes from the root to the
	current node.
	"""
	hierarchy = []
	node: Optional['ProtoNode'] = self
	while node is not None and node != root:
	hierarchy.append(attr_accessor(node))
	node = node.parent()
	return reversed(hierarchy)

	@abc.abstractmethod
	def _supports_child(self, child: 'ProtoNode') -> bool:
	"""Returns True if child is a valid child type for the current node."""


	class ProtoPackage(ProtoNode):
	"""A protobuf package."""
	def type(self) -> ProtoNode.Type:
	return ProtoNode.Type.PACKAGE

	def _supports_child(self, child: ProtoNode) -> bool:
	return True


	class ProtoEnum(ProtoNode):
	"""Representation of an enum in a .proto file."""

	# Prefix for names of values in C++ enums.
	ENUM_PREFIX: str = 'k'

	def __init__(self, name: str):
	super().__init__(name)
	self._values: List[Tuple[str, int]] = []

	def type(self) -> ProtoNode.Type:
	return ProtoNode.Type.ENUM

	def values(self) -> List[Tuple[str, int]]:
	return list(self._values)

	def add_value(self, name: str, value: int) -> None:
	name = '{}{}'.format(ProtoEnum.ENUM_PREFIX,
	ProtoMessageField.canonicalize_name(name))
	self._values.append((name, value))

	def _supports_child(self, child: ProtoNode) -> bool:
	# Enums cannot have nested children.
	return False


	class ProtoMessage(ProtoNode):
	"""Representation of a message in a .proto file."""
	def __init__(self, name: str):
	super().__init__(name)
	self._fields: List['ProtoMessageField'] = []

	def type(self) -> ProtoNode.Type:
	return ProtoNode.Type.MESSAGE

	def fields(self) -> List['ProtoMessageField']:
	return list(self._fields)

	def add_field(self, field: 'ProtoMessageField') -> None:
	self._fields.append(field)

	def _supports_child(self, child: ProtoNode) -> bool:
	return (child.type() == self.Type.ENUM
	or child.type() == self.Type.MESSAGE)


	# This class is not a node and does not appear in the proto tree.
	# Fields belong to proto messages and are processed separately.
	class ProtoMessageField:
	"""Representation of a field within a protobuf message."""
	def __init__(self,
	field_name: str,
	field_number: int,
	field_type: int,
	type_node: Optional[ProtoNode] = None,
	repeated: bool = False):
	self._name: str = self.canonicalize_name(field_name)
	self._number: int = field_number
	self._type: int = field_type
	self._type_node: Optional[ProtoNode] = type_node
	self._repeated: bool = repeated

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

	def enum_name(self) -> str:
	return '{}{}'.format(ProtoEnum.ENUM_PREFIX, self._name)

	def number(self) -> int:
	return self._number

	def type(self) -> int:
	return self._type

	def type_node(self) -> Optional[ProtoNode]:
	return self._type_node

	def is_repeated(self) -> bool:
	return self._repeated

	@staticmethod
	def canonicalize_name(field_name: str) -> str:
	"""Converts a field name to UpperCamelCase."""
	name_components = field_name.split('_')
	for i, _ in enumerate(name_components):
	name_components[i] = name_components[i].lower().capitalize()
	return ''.join(name_components)