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 RegisterMessage(self, message):55    """Registers the given message type in the local database.56 57    Calls to GetSymbol() and GetMessages() will return messages registered here.58 59    Args:60      message: A :class:`google.protobuf.message.Message` subclass (or61        instance); its descriptor will be registered.62 63    Returns:64      The provided message.65    """66 67    desc = message.DESCRIPTOR68    self._classes[desc] = message69    self.RegisterMessageDescriptor(desc)70    return message71 72  def RegisterMessageDescriptor(self, message_descriptor):73    """Registers the given message descriptor in the local database.74 75    Args:76      message_descriptor (Descriptor): the message descriptor to add.77    """78    if api_implementation.Type() == 'python':79      # pylint: disable=protected-access80      self.pool._AddDescriptor(message_descriptor)81 82  def RegisterEnumDescriptor(self, enum_descriptor):83    """Registers the given enum descriptor in the local database.84 85    Args:86      enum_descriptor (EnumDescriptor): The enum descriptor to register.87 88    Returns:89      EnumDescriptor: The provided descriptor.90    """91    if api_implementation.Type() == 'python':92      # pylint: disable=protected-access93      self.pool._AddEnumDescriptor(enum_descriptor)94    return enum_descriptor95 96  def RegisterServiceDescriptor(self, service_descriptor):97    """Registers the given service descriptor in the local database.98 99    Args:100      service_descriptor (ServiceDescriptor): the service descriptor to101        register.102    """103    if api_implementation.Type() == 'python':104      # pylint: disable=protected-access105      self.pool._AddServiceDescriptor(service_descriptor)106 107  def RegisterFileDescriptor(self, file_descriptor):108    """Registers the given file descriptor in the local database.109 110    Args:111      file_descriptor (FileDescriptor): The file descriptor to register.112    """113    if api_implementation.Type() == 'python':114      # pylint: disable=protected-access115      self.pool._InternalAddFileDescriptor(file_descriptor)116 117  def GetSymbol(self, symbol):118    """Tries to find a symbol in the local database.119 120    Currently, this method only returns message.Message instances, however, if121    may be extended in future to support other symbol types.122 123    Args:124      symbol (str): a protocol buffer symbol.125 126    Returns:127      A Python class corresponding to the symbol.128 129    Raises:130      KeyError: if the symbol could not be found.131    """132 133    return self._classes[self.pool.FindMessageTypeByName(symbol)]134 135  def GetMessages(self, files):136    # TODO: Fix the differences with MessageFactory.137    """Gets all registered messages from a specified file.138 139    Only messages already created and registered will be returned; (this is the140    case for imported _pb2 modules)141    But unlike MessageFactory, this version also returns already defined nested142    messages, but does not register any message extensions.143 144    Args:145      files (list[str]): The file names to extract messages from.146 147    Returns:148      A dictionary mapping proto names to the message classes.149 150    Raises:151      KeyError: if a file could not be found.152    """153 154    def _GetAllMessages(desc):155      """Walk a message Descriptor and recursively yields all message names."""156      yield desc157      for msg_desc in desc.nested_types:158        for nested_desc in _GetAllMessages(msg_desc):159          yield nested_desc160 161    result = {}162    for file_name in files:163      file_desc = self.pool.FindFileByName(file_name)164      for msg_desc in file_desc.message_types_by_name.values():165        for desc in _GetAllMessages(msg_desc):166          try:167            result[desc.full_name] = self._classes[desc]168          except KeyError:169            # This descriptor has no registered class, skip it.170            pass171    return result172 173 174_DEFAULT = SymbolDatabase(pool=descriptor_pool.Default())175 176 177def Default():178  """Returns the default SymbolDatabase."""179  return _DEFAULT180