python/google/net/proto2/python/public/symbol_database.py

   1 #!/usr/bin/env python
   2 #
   3 # Copyright 2007 Google Inc.
   4 #
   5 # Licensed under the Apache License, Version 2.0 (the "License");
   6 # you may not use this file except in compliance with the License.
   7 # You may obtain a copy of the License at
   8 #
   9 #     http://www.apache.org/licenses/LICENSE-2.0
  10 #
  11 # Unless required by applicable law or agreed to in writing, software
  12 # distributed under the License is distributed on an "AS IS" BASIS,
  13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14 # See the License for the specific language governing permissions and
  15 # limitations under the License.
  16 #
  17 """A database of Python protocol buffer generated symbols.
  18
  19 SymbolDatabase makes it easy to create new instances of a registered type, given
  20 only the type's protocol buffer symbol name. Once all symbols are registered,
  21 they can be accessed using either the MessageFactory interface which
  22 SymbolDatabase exposes, or the BasicDescriptorPool interface of the underlying
  23 pool.
  24
  25 Example usage:
  26
  27   db = symbol_database.SymbolDatabase()
  28
  29   # Register symbols of interest, from one or multiple files.
  30   db.RegisterFileDescriptor(my_proto_pb2.DESCRIPTOR)
  31   db.RegisterMessage(my_proto_pb2.MyMessage)
  32   db.RegisterEnumDescriptor(my_proto_pb2.MyEnum.DESCRIPTOR)
  33
  34   # The database can be used as a MessageFactory, to generate types based on
  35   # their name:
  36   types = db.GetMessages(['my_proto.proto'])
  37   my_message_instance = types['MyMessage']()
  38
  39   # The database's underlying descriptor pool can be queried, so it's not
  40   # necessary to know a type's filename to be able to generate it:
  41   filename = db.pool.FindFileContainingSymbol('MyMessage')
  42   my_message_instance = db.GetMessages([filename])['MyMessage']()
  43
  44   # This functionality is also provided directly via a convenience method:
  45   my_message_instance = db.GetSymbol('MyMessage')()
  46 """
  47
  48
  49 from google.net.proto2.python.public import basic_descriptor_pool
  50
  51
  52 class SymbolDatabase(object):
  53   """A database of Python generated symbols.
  54
  55   SymbolDatabase also models message_factory.MessageFactory.
  56
  57   The symbol database can be used to keep a global registry of all protocol
  58   buffer types used within a program.
  59   """
  60
  61   def __init__(self):
  62     """Constructor."""
  63
  64     self._symbols = {}
  65     self._symbols_by_file = {}
  66     self.pool = basic_descriptor_pool.BasicDescriptorPool()
  67
  68   def RegisterMessage(self, message):
  69     """Registers the given message type in the local database.
  70
  71     Args:
  72       message: a message.Message, to be registered.
  73
  74     Returns:
  75       The provided message.
  76     """
  77
  78     desc = message.DESCRIPTOR
  79     self._symbols[desc.full_name] = message
  80     if desc.file.name not in self._symbols_by_file:
  81       self._symbols_by_file[desc.file.name] = {}
  82     self._symbols_by_file[desc.file.name][desc.full_name] = message
  83     self.pool.AddMessage(desc)
  84     return message
  85
  86   def RegisterEnumDescriptor(self, enum_descriptor):
  87     """Registers the given enum descriptor in the local database.
  88
  89     Args:
  90       enum_descriptor: a descriptor.EnumDescriptor.
  91
  92     Returns:
  93       The provided descriptor.
  94     """
  95     self.pool.AddEnum(enum_descriptor)
  96     return enum_descriptor
  97
  98   def RegisterFileDescriptor(self, file_descriptor):
  99     """Registers the given file descriptor in the local database.
 100
 101     Args:
 102       file_descriptor: a descriptor.FileDescriptor.
 103
 104     Returns:
 105       The provided descriptor.
 106     """
 107     self.pool.AddFile(file_descriptor)
 108
 109   def GetSymbol(self, symbol):
 110     """Tries to find a symbol in the local database.
 111
 112     Currently, this method only returns message.Message instances, however, if
 113     may be extended in future to support other symbol types.
 114
 115     Args:
 116       symbol: A str, a protocol buffer symbol.
 117
 118     Returns:
 119       A Python class corresponding to the symbol.
 120
 121     Raises:
 122       KeyError: if the symbol could not be found.
 123     """
 124
 125     return self._symbols[symbol]
 126
 127   def GetPrototype(self, descriptor):
 128     """Builds a proto2 message class based on the passed in descriptor.
 129
 130     Passing a descriptor with a fully qualified name matching a previous
 131     invocation will cause the same class to be returned.
 132
 133     Args:
 134       descriptor: The descriptor to build from.
 135
 136     Returns:
 137       A class describing the passed in descriptor.
 138     """
 139
 140     return self.GetSymbol(descriptor.full_name)
 141
 142   def GetMessages(self, files):
 143     """Gets all the messages from a specified file.
 144
 145     This will find and resolve dependencies, failing if they are not registered
 146     in the symbol database.
 147
 148
 149     Args:
 150       files: The file names to extract messages from.
 151
 152     Returns:
 153       A dictionary mapping proto names to the message classes. This will include
 154       any dependent messages as well as any messages defined in the same file as
 155       a specified message.
 156
 157     Raises:
 158       KeyError: if a file could not be found.
 159     """
 160
 161     result = {}
 162     for f in files:
 163       result.update(self._symbols_by_file[f])
 164     return result
 165
 166 _DEFAULT = SymbolDatabase()
 167
 168
 169 def Default():
 170   """Returns the default SymbolDatabase."""
 171   return _DEFAULT