codekingpro/portable-devtools

1# Protocol Buffers - Google's data interchange format2# Copyright 2008 Google Inc.  All rights reserved.3#4# Use of this source code is governed by a BSD-style5# license that can be found in the LICENSE file or at6# https://developers.google.com/open-source/licenses/bsd7 8"""A database of Python protocol buffer generated symbols.9 10SymbolDatabase is the MessageFactory for messages generated at compile time,11and makes it easy to create new instances of a registered type, given only the12type's protocol buffer symbol name.13 14Example usage::15 16  db = symbol_database.SymbolDatabase()17 18  # Register symbols of interest, from one or multiple files.19  db.RegisterFileDescriptor(my_proto_pb2.DESCRIPTOR)20  db.RegisterMessage(my_proto_pb2.MyMessage)21  db.RegisterEnumDescriptor(my_proto_pb2.MyEnum.DESCRIPTOR)22 23  # The database can be used as a MessageFactory, to generate types based on24  # their name:25  types = db.GetMessages(['my_proto.proto'])26  my_message_instance = types['MyMessage']()27 28  # The database's underlying descriptor pool can be queried, so it's not29  # necessary to know a type's filename to be able to generate it:30  filename = db.pool.FindFileContainingSymbol('MyMessage')31  my_message_instance = db.GetMessages([filename])['MyMessage']()32 33  # This functionality is also provided directly via a convenience method:34  my_message_instance = db.GetSymbol('MyMessage')()35"""36 37import warnings38 39from google.protobuf.internal import api_implementation40from google.protobuf import descriptor_pool41from google.protobuf import message_factory42 43 44class SymbolDatabase():45  """A database of Python generated symbols."""46 47  # local cache of registered classes.48  _classes = {}49 50  def __init__(self, pool=None):51    """Initializes a new SymbolDatabase."""52    self.pool = pool or descriptor_pool.DescriptorPool()53 54  def GetPrototype(self, descriptor):55    warnings.warn('SymbolDatabase.GetPrototype() is deprecated. Please '56                  'use message_factory.GetMessageClass() instead. '57                  'SymbolDatabase.GetPrototype() will be removed soon.')58    return message_factory.GetMessageClass(descriptor)59 60  def CreatePrototype(self, descriptor):61    warnings.warn('Directly call CreatePrototype() is wrong. Please use '62                  'message_factory.GetMessageClass() instead. '63                  'SymbolDatabase.CreatePrototype() will be removed soon.')64    return message_factory._InternalCreateMessageClass(descriptor)65 66  def GetMessages(self, files):67    warnings.warn('SymbolDatabase.GetMessages() is deprecated. Please use '68                  'message_factory.GetMessageClassedForFiles() instead. '69                  'SymbolDatabase.GetMessages() will be removed soon.')70    return message_factory.GetMessageClassedForFiles(files, self.pool)71 72  def RegisterMessage(self, message):73    """Registers the given message type in the local database.74 75    Calls to GetSymbol() and GetMessages() will return messages registered here.76 77    Args:78      message: A :class:`google.protobuf.message.Message` subclass (or79        instance); its descriptor will be registered.80 81    Returns:82      The provided message.83    """84 85    desc = message.DESCRIPTOR86    self._classes[desc] = message87    self.RegisterMessageDescriptor(desc)88    return message89 90  def RegisterMessageDescriptor(self, message_descriptor):91    """Registers the given message descriptor in the local database.92 93    Args:94      message_descriptor (Descriptor): the message descriptor to add.95    """96    if api_implementation.Type() == 'python':97      # pylint: disable=protected-access98      self.pool._AddDescriptor(message_descriptor)99 100  def RegisterEnumDescriptor(self, enum_descriptor):101    """Registers the given enum descriptor in the local database.102 103    Args:104      enum_descriptor (EnumDescriptor): The enum descriptor to register.105 106    Returns:107      EnumDescriptor: The provided descriptor.108    """109    if api_implementation.Type() == 'python':110      # pylint: disable=protected-access111      self.pool._AddEnumDescriptor(enum_descriptor)112    return enum_descriptor113 114  def RegisterServiceDescriptor(self, service_descriptor):115    """Registers the given service descriptor in the local database.116 117    Args:118      service_descriptor (ServiceDescriptor): the service descriptor to119        register.120    """121    if api_implementation.Type() == 'python':122      # pylint: disable=protected-access123      self.pool._AddServiceDescriptor(service_descriptor)124 125  def RegisterFileDescriptor(self, file_descriptor):126    """Registers the given file descriptor in the local database.127 128    Args:129      file_descriptor (FileDescriptor): The file descriptor to register.130    """131    if api_implementation.Type() == 'python':132      # pylint: disable=protected-access133      self.pool._InternalAddFileDescriptor(file_descriptor)134 135  def GetSymbol(self, symbol):136    """Tries to find a symbol in the local database.137 138    Currently, this method only returns message.Message instances, however, if139    may be extended in future to support other symbol types.140 141    Args:142      symbol (str): a protocol buffer symbol.143 144    Returns:145      A Python class corresponding to the symbol.146 147    Raises:148      KeyError: if the symbol could not be found.149    """150 151    return self._classes[self.pool.FindMessageTypeByName(symbol)]152 153  def GetMessages(self, files):154    # TODO: Fix the differences with MessageFactory.155    """Gets all registered messages from a specified file.156 157    Only messages already created and registered will be returned; (this is the158    case for imported _pb2 modules)159    But unlike MessageFactory, this version also returns already defined nested160    messages, but does not register any message extensions.161 162    Args:163      files (list[str]): The file names to extract messages from.164 165    Returns:166      A dictionary mapping proto names to the message classes.167 168    Raises:169      KeyError: if a file could not be found.170    """171 172    def _GetAllMessages(desc):173      """Walk a message Descriptor and recursively yields all message names."""174      yield desc175      for msg_desc in desc.nested_types:176        for nested_desc in _GetAllMessages(msg_desc):177          yield nested_desc178 179    result = {}180    for file_name in files:181      file_desc = self.pool.FindFileByName(file_name)182      for msg_desc in file_desc.message_types_by_name.values():183        for desc in _GetAllMessages(msg_desc):184          try:185            result[desc.full_name] = self._classes[desc]186          except KeyError:187            # This descriptor has no registered class, skip it.188            pass189    return result190 191 192_DEFAULT = SymbolDatabase(pool=descriptor_pool.Default())193 194 195def Default():196  """Returns the default SymbolDatabase."""197  return _DEFAULT198