Correct collection packing for v3 protocol

This commit is contained in:
Tyler Hobbs
2014-06-24 18:24:11 -05:00
parent 82477e20ac
commit fce517d87f

View File

@@ -699,11 +699,12 @@ class _SimpleParameterizedType(_ParameterizedType):
raise TypeError("Received a string for a type that expects a sequence") raise TypeError("Received a string for a type that expects a sequence")
subtype, = cls.subtypes subtype, = cls.subtypes
pack = int32_pack if protocol_version >= 3 else uint16_pack
buf = io.BytesIO() buf = io.BytesIO()
buf.write(uint16_pack(len(items))) buf.write(pack(len(items)))
for item in items: for item in items:
itembytes = subtype.to_binary(item, protocol_version) itembytes = subtype.to_binary(item, protocol_version)
buf.write(uint16_pack(len(itembytes))) buf.write(pack(len(itembytes)))
buf.write(itembytes) buf.write(itembytes)
return buf.getvalue() return buf.getvalue()
@@ -758,8 +759,9 @@ class MapType(_ParameterizedType):
@classmethod @classmethod
def serialize_safe(cls, themap, protocol_version): def serialize_safe(cls, themap, protocol_version):
subkeytype, subvaltype = cls.subtypes subkeytype, subvaltype = cls.subtypes
pack = int32_pack if protocol_version >= 3 else uint16_pack
buf = io.BytesIO() buf = io.BytesIO()
buf.write(uint16_pack(len(themap))) buf.write(pack(len(themap)))
try: try:
items = six.iteritems(themap) items = six.iteritems(themap)
except AttributeError: except AttributeError:
@@ -767,9 +769,9 @@ class MapType(_ParameterizedType):
for key, val in items: for key, val in items:
keybytes = subkeytype.to_binary(key, protocol_version) keybytes = subkeytype.to_binary(key, protocol_version)
valbytes = subvaltype.to_binary(val, protocol_version) valbytes = subvaltype.to_binary(val, protocol_version)
buf.write(uint16_pack(len(keybytes))) buf.write(pack(len(keybytes)))
buf.write(keybytes) buf.write(keybytes)
buf.write(uint16_pack(len(valbytes))) buf.write(pack(len(valbytes)))
buf.write(valbytes) buf.write(valbytes)
return buf.getvalue() return buf.getvalue()