• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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.internal import api_implementation
62from google.protobuf import descriptor_pool
63from google.protobuf import message_factory
64
65
66class SymbolDatabase(message_factory.MessageFactory):
67  """A database of Python generated symbols."""
68
69  def RegisterMessage(self, message):
70    """Registers the given message type in the local database.
71
72    Calls to GetSymbol() and GetMessages() will return messages registered here.
73
74    Args:
75      message: A :class:`google.protobuf.message.Message` subclass (or
76        instance); its descriptor will be registered.
77
78    Returns:
79      The provided message.
80    """
81
82    desc = message.DESCRIPTOR
83    self._classes[desc] = message
84    self.RegisterMessageDescriptor(desc)
85    return message
86
87  def RegisterMessageDescriptor(self, message_descriptor):
88    """Registers the given message descriptor in the local database.
89
90    Args:
91      message_descriptor (Descriptor): the message descriptor to add.
92    """
93    if api_implementation.Type() == 'python':
94      # pylint: disable=protected-access
95      self.pool._AddDescriptor(message_descriptor)
96
97  def RegisterEnumDescriptor(self, enum_descriptor):
98    """Registers the given enum descriptor in the local database.
99
100    Args:
101      enum_descriptor (EnumDescriptor): The enum descriptor to register.
102
103    Returns:
104      EnumDescriptor: The provided descriptor.
105    """
106    if api_implementation.Type() == 'python':
107      # pylint: disable=protected-access
108      self.pool._AddEnumDescriptor(enum_descriptor)
109    return enum_descriptor
110
111  def RegisterServiceDescriptor(self, service_descriptor):
112    """Registers the given service descriptor in the local database.
113
114    Args:
115      service_descriptor (ServiceDescriptor): the service descriptor to
116        register.
117    """
118    if api_implementation.Type() == 'python':
119      # pylint: disable=protected-access
120      self.pool._AddServiceDescriptor(service_descriptor)
121
122  def RegisterFileDescriptor(self, file_descriptor):
123    """Registers the given file descriptor in the local database.
124
125    Args:
126      file_descriptor (FileDescriptor): The file descriptor to register.
127    """
128    if api_implementation.Type() == 'python':
129      # pylint: disable=protected-access
130      self.pool._InternalAddFileDescriptor(file_descriptor)
131
132  def GetSymbol(self, symbol):
133    """Tries to find a symbol in the local database.
134
135    Currently, this method only returns message.Message instances, however, if
136    may be extended in future to support other symbol types.
137
138    Args:
139      symbol (str): a protocol buffer symbol.
140
141    Returns:
142      A Python class corresponding to the symbol.
143
144    Raises:
145      KeyError: if the symbol could not be found.
146    """
147
148    return self._classes[self.pool.FindMessageTypeByName(symbol)]
149
150  def GetMessages(self, files):
151    # TODO(amauryfa): Fix the differences with MessageFactory.
152    """Gets all registered messages from a specified file.
153
154    Only messages already created and registered will be returned; (this is the
155    case for imported _pb2 modules)
156    But unlike MessageFactory, this version also returns already defined nested
157    messages, but does not register any message extensions.
158
159    Args:
160      files (list[str]): The file names to extract messages from.
161
162    Returns:
163      A dictionary mapping proto names to the message classes.
164
165    Raises:
166      KeyError: if a file could not be found.
167    """
168
169    def _GetAllMessages(desc):
170      """Walk a message Descriptor and recursively yields all message names."""
171      yield desc
172      for msg_desc in desc.nested_types:
173        for nested_desc in _GetAllMessages(msg_desc):
174          yield nested_desc
175
176    result = {}
177    for file_name in files:
178      file_desc = self.pool.FindFileByName(file_name)
179      for msg_desc in file_desc.message_types_by_name.values():
180        for desc in _GetAllMessages(msg_desc):
181          try:
182            result[desc.full_name] = self._classes[desc]
183          except KeyError:
184            # This descriptor has no registered class, skip it.
185            pass
186    return result
187
188
189_DEFAULT = SymbolDatabase(pool=descriptor_pool.Default())
190
191
192def Default():
193  """Returns the default SymbolDatabase."""
194  return _DEFAULT
195