1# Protocol Buffers - Google's data interchange format
2# Copyright 2008 Google Inc.  All rights reserved.
3# https://developers.google.com/protocol-buffers/
4#
5# Redistribution and use in source and binary forms, with or without
6# modification, are permitted provided that the following conditions are
7# met:
8#
9#     * Redistributions of source code must retain the above copyright
10# notice, this list of conditions and the following disclaimer.
11#     * Redistributions in binary form must reproduce the above
12# copyright notice, this list of conditions and the following disclaimer
13# in the documentation and/or other materials provided with the
14# distribution.
15#     * Neither the name of Google Inc. nor the names of its
16# contributors may be used to endorse or promote products derived from
17# this software without specific prior written permission.
18#
19# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
20# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
21# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
22# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
23# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
24# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
25# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
26# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
27# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
28# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30
31"""A database of Python protocol buffer generated symbols.
32
33SymbolDatabase is the MessageFactory for messages generated at compile time,
34and makes it easy to create new instances of a registered type, given only the
35type's protocol buffer symbol name.
36
37Example usage:
38
39  db = symbol_database.SymbolDatabase()
40
41  # Register symbols of interest, from one or multiple files.
42  db.RegisterFileDescriptor(my_proto_pb2.DESCRIPTOR)
43  db.RegisterMessage(my_proto_pb2.MyMessage)
44  db.RegisterEnumDescriptor(my_proto_pb2.MyEnum.DESCRIPTOR)
45
46  # The database can be used as a MessageFactory, to generate types based on
47  # their name:
48  types = db.GetMessages(['my_proto.proto'])
49  my_message_instance = types['MyMessage']()
50
51  # The database's underlying descriptor pool can be queried, so it's not
52  # necessary to know a type's filename to be able to generate it:
53  filename = db.pool.FindFileContainingSymbol('MyMessage')
54  my_message_instance = db.GetMessages([filename])['MyMessage']()
55
56  # This functionality is also provided directly via a convenience method:
57  my_message_instance = db.GetSymbol('MyMessage')()
58"""
59
60
61from google.protobuf import descriptor_pool
62from google.protobuf import message_factory
63
64
65class SymbolDatabase(message_factory.MessageFactory):
66  """A database of Python generated symbols."""
67
68  def RegisterMessage(self, message):
69    """Registers the given message type in the local database.
70
71    Calls to GetSymbol() and GetMessages() will return messages registered here.
72
73    Args:
74      message: a message.Message, to be registered.
75
76    Returns:
77      The provided message.
78    """
79
80    desc = message.DESCRIPTOR
81    self._classes[desc.full_name] = message
82    self.pool.AddDescriptor(desc)
83    return message
84
85  def RegisterEnumDescriptor(self, enum_descriptor):
86    """Registers the given enum descriptor in the local database.
87
88    Args:
89      enum_descriptor: a descriptor.EnumDescriptor.
90
91    Returns:
92      The provided descriptor.
93    """
94    self.pool.AddEnumDescriptor(enum_descriptor)
95    return enum_descriptor
96
97  def RegisterFileDescriptor(self, file_descriptor):
98    """Registers the given file descriptor in the local database.
99
100    Args:
101      file_descriptor: a descriptor.FileDescriptor.
102
103    Returns:
104      The provided descriptor.
105    """
106    self.pool.AddFileDescriptor(file_descriptor)
107
108  def GetSymbol(self, symbol):
109    """Tries to find a symbol in the local database.
110
111    Currently, this method only returns message.Message instances, however, if
112    may be extended in future to support other symbol types.
113
114    Args:
115      symbol: A str, a protocol buffer symbol.
116
117    Returns:
118      A Python class corresponding to the symbol.
119
120    Raises:
121      KeyError: if the symbol could not be found.
122    """
123
124    return self._classes[symbol]
125
126  def GetMessages(self, files):
127    # TODO(amauryfa): Fix the differences with MessageFactory.
128    """Gets all registered messages from a specified file.
129
130    Only messages already created and registered will be returned; (this is the
131    case for imported _pb2 modules)
132    But unlike MessageFactory, this version also returns nested messages.
133
134    Args:
135      files: The file names to extract messages from.
136
137    Returns:
138      A dictionary mapping proto names to the message classes.
139
140    Raises:
141      KeyError: if a file could not be found.
142    """
143
144    def _GetAllMessageNames(desc):
145      """Walk a message Descriptor and recursively yields all message names."""
146      yield desc.full_name
147      for msg_desc in desc.nested_types:
148        for full_name in _GetAllMessageNames(msg_desc):
149          yield full_name
150
151    result = {}
152    for file_name in files:
153      file_desc = self.pool.FindFileByName(file_name)
154      for msg_desc in file_desc.message_types_by_name.values():
155        for full_name in _GetAllMessageNames(msg_desc):
156          try:
157            result[full_name] = self._classes[full_name]
158          except KeyError:
159            # This descriptor has no registered class, skip it.
160            pass
161    return result
162
163
164_DEFAULT = SymbolDatabase(pool=descriptor_pool.Default())
165
166
167def Default():
168  """Returns the default SymbolDatabase."""
169  return _DEFAULT
170