512 lines
18 KiB
Python
512 lines
18 KiB
Python
# Copyright 2016 Google LLC. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
# Copyright 2005 Google LLC.
|
|
|
|
"""This module contains the MessageSet class, which is a special kind of
|
|
protocol message which can contain other protocol messages without knowing
|
|
their types. See the class's doc string for more information."""
|
|
|
|
# When a MessageSet is first parsed, it contains a set of strings. It
|
|
# can't parse these to the final message types until it knows what type they
|
|
# are. When the user first requests a particular message type, the actual
|
|
# class implementing the message type is captured for the MessageSet to
|
|
# use. The string is then replaced with its fully-parsed equivalent. Thus,
|
|
# each message is only parsed on-demand, the first time it is requested.
|
|
|
|
from googlecloudsdk.appengine.proto import ProtocolBuffer
|
|
from googlecloudsdk.core import log
|
|
|
|
# Degrade gracefully for users of the :pure_pyproto target who depend on
|
|
# some MessageSet functionality, but cannot include this C++ extension.
|
|
try:
|
|
from googlecloudsdk.appengine.proto import _net_proto___parse__python
|
|
except ImportError:
|
|
_net_proto___parse__python = None
|
|
|
|
# Tags byte values used in the wire format of MessageSet.
|
|
TAG_BEGIN_ITEM_GROUP = 11
|
|
TAG_END_ITEM_GROUP = 12
|
|
TAG_TYPE_ID = 16
|
|
TAG_MESSAGE = 26
|
|
|
|
class Item:
|
|
"""This class is used internally by MessageSet.
|
|
|
|
This class represents a particular message in the set. The message may be
|
|
either parsed or unparsed. The point of this class is to provide a nice
|
|
wrapper for all of the standard protocol buffer operations so that the
|
|
caller does not have to check if the message is parsed or not."""
|
|
|
|
def __init__(self, message, message_class=None):
|
|
self.message = message
|
|
self.message_class = message_class
|
|
|
|
def SetToDefaultInstance(self, message_class):
|
|
self.message = message_class()
|
|
self.message_class = message_class
|
|
|
|
def Parse(self, message_class):
|
|
"""Upgrades the Item to a parsed one, returning true if successful."""
|
|
|
|
if self.message_class is not None:
|
|
return 1
|
|
|
|
try:
|
|
message_obj = message_class()
|
|
message_obj.MergePartialFromString(self.message)
|
|
self.message = message_obj
|
|
self.message_class = message_class
|
|
return 1
|
|
except ProtocolBuffer.ProtocolBufferDecodeError:
|
|
logging.warn("Parse error in message inside MessageSet. Tried "
|
|
"to parse as: " + message_class.__name__)
|
|
return 0
|
|
|
|
def MergeFrom(self, other):
|
|
"""Merges two items.
|
|
|
|
If one item is parsed and one is not, an attempt will be made to parse
|
|
the unparsed one."""
|
|
|
|
if self.message_class is not None:
|
|
if other.Parse(self.message_class):
|
|
self.message.MergeFrom(other.message)
|
|
# Pretend other was empty if it didn't parse, since callers to MergeFrom
|
|
# don't expect to handle errors.
|
|
|
|
elif other.message_class is not None:
|
|
if not self.Parse(other.message_class):
|
|
# Pretend self was empty if it didn't parse.
|
|
self.message = other.message_class()
|
|
self.message_class = other.message_class
|
|
self.message.MergeFrom(other.message)
|
|
|
|
else:
|
|
self.message += other.message
|
|
|
|
def Copy(self):
|
|
"""Make a deep copy of the Item."""
|
|
|
|
if self.message_class is None:
|
|
return Item(self.message)
|
|
else:
|
|
new_message = self.message_class()
|
|
new_message.CopyFrom(self.message)
|
|
return Item(new_message, self.message_class)
|
|
|
|
def Equals(self, other):
|
|
"""Check if two items are equal.
|
|
|
|
If one message is parsed and the other is not, an attempt will be made
|
|
to parse the unparsed one."""
|
|
|
|
if self.message_class is not None:
|
|
if not other.Parse(self.message_class): return 0
|
|
return self.message.Equals(other.message)
|
|
|
|
elif other.message_class is not None:
|
|
if not self.Parse(other.message_class): return 0
|
|
return self.message.Equals(other.message)
|
|
|
|
else:
|
|
return self.message == other.message
|
|
|
|
def IsInitialized(self, debug_strs=None):
|
|
"""Calls IsInitialized on the contained message if it is parsed, or just
|
|
returns true otherwise."""
|
|
|
|
if self.message_class is None:
|
|
return 1
|
|
else:
|
|
return self.message.IsInitialized(debug_strs)
|
|
|
|
def ByteSize(self, pb, type_id):
|
|
"""Returns the encoded size of this item."""
|
|
|
|
message_length = 0
|
|
if self.message_class is None:
|
|
message_length = len(self.message)
|
|
else:
|
|
message_length = self.message.ByteSize()
|
|
|
|
# Extra two bytes is tags for type_id and message.
|
|
return pb.lengthString(message_length) + pb.lengthVarInt64(type_id) + 2
|
|
|
|
def ByteSizePartial(self, pb, type_id):
|
|
"""Returns the encoded size of this item, not counting missing optional."""
|
|
|
|
message_length = 0
|
|
if self.message_class is None:
|
|
message_length = len(self.message)
|
|
else:
|
|
message_length = self.message.ByteSizePartial()
|
|
|
|
# Extra two bytes is tags for type_id and message.
|
|
return pb.lengthString(message_length) + pb.lengthVarInt64(type_id) + 2
|
|
|
|
def OutputUnchecked(self, out, type_id):
|
|
"""Writes the item to the decoder "out", given its type ID."""
|
|
|
|
out.putVarInt32(TAG_TYPE_ID)
|
|
# Type IDs are defined as being 32-bit as of this writing, but could
|
|
# become 64-bit in the future. Since putVarInt32() is just a wrapper
|
|
# for putVarUint64() with added bounds checking, we might as well use
|
|
# the 64-bit version now.
|
|
out.putVarUint64(type_id)
|
|
out.putVarInt32(TAG_MESSAGE)
|
|
if self.message_class is None:
|
|
out.putPrefixedString(self.message)
|
|
else:
|
|
out.putVarInt32(self.message.ByteSize())
|
|
self.message.OutputUnchecked(out)
|
|
|
|
def OutputPartial(self, out, type_id):
|
|
"""Writes the item to the decoder "out", given its type ID.
|
|
Does not assume required fields are set."""
|
|
|
|
out.putVarInt32(TAG_TYPE_ID)
|
|
# Type IDs are defined as being 32-bit as of this writing, but could
|
|
# become 64-bit in the future. Since putVarInt32() is just a wrapper
|
|
# for putVarUint64() with added bounds checking, we might as well use
|
|
# the 64-bit version now.
|
|
out.putVarUint64(type_id)
|
|
out.putVarInt32(TAG_MESSAGE)
|
|
if self.message_class is None:
|
|
out.putPrefixedString(self.message)
|
|
else:
|
|
out.putVarInt32(self.message.ByteSizePartial())
|
|
self.message.OutputPartial(out)
|
|
|
|
def Decode(decoder):
|
|
"""Decodes a type_id and message buffer from the decoder. (static)
|
|
|
|
Returns a (type_id, message) tuple."""
|
|
|
|
type_id = 0
|
|
message = None
|
|
while 1:
|
|
tag = decoder.getVarInt32()
|
|
if tag == TAG_END_ITEM_GROUP:
|
|
break
|
|
if tag == TAG_TYPE_ID:
|
|
# Type IDs are defined as being 32-bit as of this writing, but could
|
|
# become 64-bit in the future. Since getVarInt32() is just a wrapper
|
|
# for getVarUint64() with added bounds hcecking, we might as well use
|
|
# the 64-bit version now.
|
|
type_id = decoder.getVarUint64()
|
|
continue
|
|
if tag == TAG_MESSAGE:
|
|
message = decoder.getPrefixedString()
|
|
continue
|
|
# tag 0 is special: it's used to indicate an error.
|
|
# so if we see it we raise an exception.
|
|
if tag == 0: raise ProtocolBuffer.ProtocolBufferDecodeError
|
|
decoder.skipData(tag)
|
|
|
|
if type_id == 0 or message is None:
|
|
raise ProtocolBuffer.ProtocolBufferDecodeError
|
|
return (type_id, message)
|
|
Decode = staticmethod(Decode)
|
|
|
|
|
|
class MessageSet(ProtocolBuffer.ProtocolMessage):
|
|
"""A protocol message which contains other protocol messages.
|
|
|
|
This class is a specially-crafted ProtocolMessage which represents a
|
|
container storing other protocol messages. The contained messages can be
|
|
of any protocol message type which has been assigned a unique type
|
|
identifier. No two messages in the MessageSet may have the same type,
|
|
but otherwise there are no restrictions on what you can put in it. Accessing
|
|
the stored messages involves passing the class objects representing the
|
|
types you are looking for:
|
|
assert myMessageSet.has(MyMessageType)
|
|
message = myMessageSet.get(MyMessageType)
|
|
message = myMessageSet.mutable(MyMessageType)
|
|
myMessageSet.remove(MyMessageType)
|
|
|
|
Message types designed to be stored in MessageSets must have unique
|
|
32-bit type identifiers. Give your message type an identifier like so:
|
|
parsed message MyMessageType {
|
|
enum TypeId {MESSAGE_TYPE_ID = 12345678};
|
|
To insure that your type ID is unique, use one of your perforce change
|
|
numbers. Just make sure you only use your own numbers and that you don't
|
|
use the same one twice.
|
|
|
|
The wire format of a MessageSet looks like this:
|
|
parsed message MessageSet {
|
|
repeated group Item = 1 {
|
|
required fixed32 type_id = 2;
|
|
required message<RawMessage> message = 3;
|
|
};
|
|
};
|
|
The MessageSet class provides a special interface around this format for
|
|
the sake of ease-of-use and type safety.
|
|
|
|
See message_set_unittest.proto and message_set_py_unittest.py for examples.
|
|
"""
|
|
|
|
def __init__(self, contents=None):
|
|
"""Construct a new MessageSet, with optional starting contents
|
|
in binary protocol buffer format."""
|
|
self.items = dict()
|
|
if contents is not None: self.MergeFromString(contents)
|
|
|
|
# -------------------------------------------------------------------
|
|
# MessageSet standard interface.
|
|
|
|
def get(self, message_class):
|
|
"""Gets a message of the given type from the set.
|
|
|
|
If the MessageSet contains no message of that type, a default instance
|
|
of the class is returned. This is done to match the behavior of the
|
|
accessors normally generated for an optional field of a protocol message.
|
|
This makes it easier to transition from a long list of optional fields
|
|
to a MessageSet.
|
|
|
|
The returned message should not be modified.
|
|
"""
|
|
|
|
if message_class.MESSAGE_TYPE_ID not in self.items:
|
|
return message_class()
|
|
item = self.items[message_class.MESSAGE_TYPE_ID]
|
|
if item.Parse(message_class):
|
|
return item.message
|
|
else:
|
|
return message_class()
|
|
|
|
def mutable(self, message_class):
|
|
"""Gets a mutable reference to a message of the given type in the set.
|
|
|
|
If the MessageSet contains no message of that type, one is created and
|
|
added to the set.
|
|
"""
|
|
|
|
if message_class.MESSAGE_TYPE_ID not in self.items:
|
|
message = message_class()
|
|
self.items[message_class.MESSAGE_TYPE_ID] = Item(message, message_class)
|
|
return message
|
|
item = self.items[message_class.MESSAGE_TYPE_ID]
|
|
if not item.Parse(message_class):
|
|
item.SetToDefaultInstance(message_class)
|
|
return item.message
|
|
|
|
def has(self, message_class):
|
|
"""Checks if the set contains a message of the given type."""
|
|
|
|
if message_class.MESSAGE_TYPE_ID not in self.items:
|
|
return 0
|
|
item = self.items[message_class.MESSAGE_TYPE_ID]
|
|
return item.Parse(message_class)
|
|
|
|
def has_unparsed(self, message_class):
|
|
"""Checks if the set contains an unparsed message of the given type.
|
|
|
|
This differs from has() when the set contains a message of the given type
|
|
with a parse error. has() will return false when this is the case, but
|
|
has_unparsed() will return true. This is only useful for error checking.
|
|
"""
|
|
return message_class.MESSAGE_TYPE_ID in self.items
|
|
|
|
def GetTypeIds(self):
|
|
"""Return a list of all type ids in the set.
|
|
|
|
Returns:
|
|
[ cls1.MESSAGE_TYPE_ID, ... ] for each cls in the set. The returned
|
|
list does not contain duplicates.
|
|
"""
|
|
return self.items.keys()
|
|
|
|
def NumMessages(self):
|
|
"""Return the number of messages in the set. For any set the following
|
|
invariant holds:
|
|
set.NumMessages() == len(set.GetTypeIds())
|
|
|
|
Returns:
|
|
number of messages in the set
|
|
"""
|
|
return len(self.items)
|
|
|
|
def remove(self, message_class):
|
|
"""Removes any message of the given type from the set."""
|
|
if message_class.MESSAGE_TYPE_ID in self.items:
|
|
del self.items[message_class.MESSAGE_TYPE_ID]
|
|
|
|
# -------------------------------------------------------------------
|
|
# Python container methods.
|
|
|
|
def __getitem__(self, message_class):
|
|
if message_class.MESSAGE_TYPE_ID not in self.items:
|
|
raise KeyError(message_class)
|
|
item = self.items[message_class.MESSAGE_TYPE_ID]
|
|
if item.Parse(message_class):
|
|
return item.message
|
|
else:
|
|
raise KeyError(message_class)
|
|
|
|
def __setitem__(self, message_class, message):
|
|
self.items[message_class.MESSAGE_TYPE_ID] = Item(message, message_class)
|
|
|
|
def __contains__(self, message_class):
|
|
return self.has(message_class)
|
|
|
|
def __delitem__(self, message_class):
|
|
self.remove(message_class)
|
|
|
|
def __len__(self):
|
|
return len(self.items)
|
|
|
|
# -------------------------------------------------------------------
|
|
# Standard ProtocolMessage interface.
|
|
|
|
def MergeFrom(self, other):
|
|
"""Merges the messages from MessageSet 'other' into this set.
|
|
|
|
If both sets contain messages of matching types, those messages will be
|
|
individually merged by type.
|
|
"""
|
|
|
|
assert other is not self
|
|
|
|
for (type_id, item) in other.items.items():
|
|
if type_id in self.items:
|
|
self.items[type_id].MergeFrom(item)
|
|
else:
|
|
self.items[type_id] = item.Copy()
|
|
|
|
def Equals(self, other):
|
|
"""Checks if two MessageSets are equal."""
|
|
if other is self: return 1
|
|
if len(self.items) != len(other.items): return 0
|
|
|
|
for (type_id, item) in other.items.items():
|
|
if type_id not in self.items: return 0
|
|
if not self.items[type_id].Equals(item): return 0
|
|
|
|
return 1
|
|
|
|
def __eq__(self, other):
|
|
return ((other is not None)
|
|
and (other.__class__ == self.__class__)
|
|
and self.Equals(other))
|
|
|
|
def __ne__(self, other):
|
|
return not (self == other)
|
|
|
|
def IsInitialized(self, debug_strs=None):
|
|
"""Checks if all messages in this set have had all of their required fields
|
|
set."""
|
|
|
|
initialized = 1
|
|
for item in self.items.values():
|
|
if not item.IsInitialized(debug_strs):
|
|
initialized = 0
|
|
return initialized
|
|
|
|
def ByteSize(self):
|
|
"""Gets the byte size of a protocol buffer representing this MessageSet."""
|
|
n = 2 * len(self.items) # 2 bytes for group start and end tags
|
|
for (type_id, item) in self.items.items():
|
|
n += item.ByteSize(self, type_id)
|
|
return n
|
|
|
|
def ByteSizePartial(self):
|
|
"""Gets the byte size of a protocol buffer representing this MessageSet.
|
|
Does not count missing required fields."""
|
|
n = 2 * len(self.items) # 2 bytes for group start and end tags
|
|
for (type_id, item) in self.items.items():
|
|
n += item.ByteSizePartial(self, type_id)
|
|
return n
|
|
|
|
def Clear(self):
|
|
"""Removes all messages from the set."""
|
|
self.items = dict()
|
|
|
|
def OutputUnchecked(self, out):
|
|
"""Writes the MessageSet to the encoder 'out'."""
|
|
for (type_id, item) in self.items.items():
|
|
out.putVarInt32(TAG_BEGIN_ITEM_GROUP)
|
|
item.OutputUnchecked(out, type_id)
|
|
out.putVarInt32(TAG_END_ITEM_GROUP)
|
|
|
|
def OutputPartial(self, out):
|
|
"""Writes the MessageSet to the encoder 'out'.
|
|
Does not assume required fields are set."""
|
|
for (type_id, item) in self.items.items():
|
|
out.putVarInt32(TAG_BEGIN_ITEM_GROUP)
|
|
item.OutputPartial(out, type_id)
|
|
out.putVarInt32(TAG_END_ITEM_GROUP)
|
|
|
|
def TryMerge(self, decoder):
|
|
"""Attempts to decode a MessageSet from the decoder 'd' and merge it
|
|
with this one."""
|
|
while decoder.avail() > 0:
|
|
tag = decoder.getVarInt32()
|
|
if tag == TAG_BEGIN_ITEM_GROUP:
|
|
(type_id, message) = Item.Decode(decoder)
|
|
if type_id in self.items:
|
|
self.items[type_id].MergeFrom(Item(message))
|
|
else:
|
|
self.items[type_id] = Item(message)
|
|
continue
|
|
# tag 0 is special: it's used to indicate an error.
|
|
# so if we see it we raise an exception.
|
|
if (tag == 0): raise ProtocolBuffer.ProtocolBufferDecodeError
|
|
decoder.skipData(tag)
|
|
|
|
# ASCII serialization. Only available through C++ extension.
|
|
def _CToASCII(self, output_format):
|
|
if _net_proto___parse__python is None:
|
|
return ProtocolBuffer.ProtocolMessage._CToASCII(self, output_format)
|
|
else:
|
|
return _net_proto___parse__python.ToASCII(
|
|
self, "MessageSetInternal", output_format)
|
|
|
|
def ParseASCII(self, s):
|
|
if _net_proto___parse__python is None:
|
|
ProtocolBuffer.ProtocolMessage.ParseASCII(self, s)
|
|
else:
|
|
_net_proto___parse__python.ParseASCII(self, "MessageSetInternal", s)
|
|
|
|
def ParseASCIIIgnoreUnknown(self, s):
|
|
if _net_proto___parse__python is None:
|
|
ProtocolBuffer.ProtocolMessage.ParseASCIIIgnoreUnknown(self, s)
|
|
else:
|
|
_net_proto___parse__python.ParseASCIIIgnoreUnknown(
|
|
self, "MessageSetInternal", s)
|
|
|
|
# NOTE(user): Use ToASCII() rather than __str__() for ASCII serialization.
|
|
# This legacy implementation has several quirks:
|
|
# - Does not qualify message names by package, so it will print "Message"
|
|
# from package foo.bar as [Message] rather than [foo.bar.Message].
|
|
# - Only knows about messages previously seen by its class instance, and
|
|
# may print byte counts in place of message data.
|
|
def __str__(self, prefix="", printElemNumber=0):
|
|
text = ""
|
|
for (type_id, item) in self.items.items():
|
|
if item.message_class is None:
|
|
text += "%s[%d] <\n" % (prefix, type_id)
|
|
text += "%s (%d bytes)\n" % (prefix, len(item.message))
|
|
text += "%s>\n" % prefix
|
|
else:
|
|
text += "%s[%s] <\n" % (prefix, item.message_class.__name__)
|
|
text += item.message.__str__(prefix + " ", printElemNumber)
|
|
text += "%s>\n" % prefix
|
|
return text
|
|
|
|
# To allow MessageSet to be used as an extension type. (b/6401336)
|
|
_PROTO_DESCRIPTOR_NAME = 'MessageSet'
|
|
|
|
__all__ = ['MessageSet']
|