mapfield.py

from protorpc import messages


class MultiMessage(messages.Message):
  """Has a field for each basic value type.
  """
  boolean = messages.BooleanField(1)
  integer = messages.IntegerField(2)
  float = messages.FloatField(3)
  string = messages.StringField(4)


class MultiField(messages.MessageField):
  """Holds either an integer, float, boolean or string.

  This is useful for when we don't know what the type of the value will be.
  """

  message_type = MultiMessage

  @util.positional(3)
  def __init__(self, number, **kwargs):
    super(MultiField, self).__init__(self.message_type, number, **kwargs)

  @property
  def type(self):
    return object

  def value_from_message(self, message):
    message = super(MultiField, self).value_from_message(message)
    if message.boolean is not None:
      return message.boolean
    if message.integer is not None:
      return message.integer
    if message.float is not None:
      return message.float
    if message.integer is not None:
      return message.string

  def value_to_message(self, value):
    value = super(MultiField, self).value_to_message(value)
    message = MultiMessage()
    if isinstance(value, bool):
      message.boolean = value
    elif isinstance(value, int):
      message.integer = value
    elif isinstance(value, float):
      message.float = value
    elif isinstance(value, basestring):
      message.string = value
    return message


class MultiMessageCollection(messages.Message):

  multi_messages = messages.MessageField(MultiMessage, 1, repeated=True)


class MapField(messages.MessageField):
  """A key-value mapping as a message field.
  """

  message_type = MultiMessageCollection

  def __init__(self, number, **kwargs):
    super(MapField, self).__init__(self.message_type, number, **kwargs)

  @property
  def type(self):
    return dict

  def value_from_message(self, message):
    """Takes a pair of values from a list of messages and makes a dict.
    """
    message = super(MapField, self).value_from_message(message)
    D = {}
    for i in range(0, len(message.multi_messages), 2):
      key_message = message.multi_messages[i]
      value_message = message.multi_messages[i+1]
      if value_message.boolean is not None:
        value = value_message.boolean
      elif value_message.integer is not None:
        value = value_message.integer
      elif value_message.float is not None:
        value = value_message.float
      elif value_message.integer is not None:
        value = value_message.string
      D[key_message.string] = value
    return D

  def value_to_message(self, value):
    """Converts a dict or list of 2-tuples to a list of messages.
    """
    value = super(MapField, self).value_to_message(value)
    messages = []
    for key,value in dict(value).iteritems():
      key_message = MultiMessage(string=key)
      value_message = MultiMessage()
      if isinstance(value, bool):
        value_message.boolean = value
      elif isinstance(value, int):
        value_message.integer = value
      elif isinstance(value, float):
        value_message.float = value
      elif isinstance(value, basestring):
        value_message.string = value
      messages.append(key_message)
      messages.append(value_message)
    return MultiMessageCollection(multi_messages=messages)