Use polymorphism for OrderedMap type when key type is known

PYTHON-231 peer review

Author: aholmberg

cassandra/cqltypes.py

@@ -48,7 +48,7 @@
                                int32_pack, int32_unpack, int64_pack, int64_unpack,
                                float_pack, float_unpack, double_pack, double_unpack,
                                varint_pack, varint_unpack)
-from cassandra.util import OrderedMap, sortedset, Time
+from cassandra.util import OrderedMapSerializedKey, sortedset, Time
 
 apache_cassandra_type_prefix = 'org.apache.cassandra.db.marshal.'
 
@@ -794,8 +794,7 @@ def deserialize_safe(cls, byts, protocol_version):
             length = 2
         numelements = unpack(byts[:length])
         p = length
-        themap = OrderedMap()
-        themap._set_key_type(key_type, protocol_version)
+        themap = OrderedMapSerializedKey(key_type, protocol_version)
         for _ in range(numelements):
             key_len = unpack(byts[p:p + length])
             p += length

cassandra/util.py

@@ -594,13 +594,6 @@ class OrderedMap(Mapping):
 
     '''
 
-    cass_key_type = None
-    protocol_version = None
-    '''
-    Set when deserializing a map from the server, when the key type is known.
-    This avoids re-serializing, and also helps normalize lookups on text types.
-    '''
-
     def __init__(self, *args, **kwargs):
         if len(args) > 1:
             raise TypeError('expected at most 1 arguments, got %d' % len(args))
@@ -628,18 +621,6 @@ def _insert(self, key, value):
             self._items.append((key, value))
             self._index[flat_key] = len(self._items) - 1
 
-    def _set_key_type(self, cass_type, protocol_version):
-        self.cass_key_type = cass_type
-        self.protocol_version = protocol_version
-
-    def _insert_unchecked(self, key, flat_key, value):
-        '''
-        Used when building from server response.
-        cass_key_type must be set in order for lookups to work later
-        '''
-        self._items.append((key, value))
-        self._index[flat_key] = len(self._items) - 1
-
     def __getitem__(self, key):
         try:
             index = self._index[self._serialize_key(key)]
@@ -675,7 +656,22 @@ def __str__(self):
         return '{%s}' % ', '.join("%s: %s" % (k, v) for k, v in self._items)
 
     def _serialize_key(self, key):
-        return self.cass_key_type.serialize(key, self.protocol_version) if self.cass_key_type else cPickle.dumps(key)
+        return cPickle.dumps(key)
+
+
+class OrderedMapSerializedKey(OrderedMap):
+
+    def __init__(self, cass_type, protocol_version):
+        super(OrderedMapSerializedKey, self).__init__()
+        self.cass_key_type = cass_type
+        self.protocol_version = protocol_version
+
+    def _insert_unchecked(self, key, flat_key, value):
+        self._items.append((key, value))
+        self._index[flat_key] = len(self._items) - 1
+
+    def _serialize_key(self, key):
+        return self.cass_key_type.serialize(key, self.protocol_version)
 
 
 import datetime

tests/unit/test_marshalling.py

@@ -23,8 +23,8 @@
 from decimal import Decimal
 from uuid import UUID
 
-from cassandra.cqltypes import lookup_casstype
-from cassandra.util import OrderedMap, sortedset, Time
+from cassandra.cqltypes import lookup_casstype, DecimalType, UTF8Type
+from cassandra.util import OrderedMap, OrderedMapSerializedKey, sortedset, Time
 
 marshalled_value_pairs = (
     # binary form, type, python native type
@@ -75,7 +75,7 @@
     (b'', 'MapType(AsciiType, BooleanType)', None),
     (b'', 'ListType(FloatType)', None),
     (b'', 'SetType(LongType)', None),
-    (b'\x00\x00', 'MapType(DecimalType, BooleanType)', OrderedMap()),
+    (b'\x00\x00', 'MapType(DecimalType, BooleanType)', OrderedMapSerializedKey(DecimalType, 0)),
     (b'\x00\x00', 'ListType(FloatType)', []),
     (b'\x00\x00', 'SetType(IntegerType)', sortedset()),
     (b'\x00\x01\x00\x10\xafYC\xa3\xea<\x11\xe1\xabc\xc4,\x03"y\xf0', 'ListType(TimeUUIDType)', [UUID(bytes=b'\xafYC\xa3\xea<\x11\xe1\xabc\xc4,\x03"y\xf0')]),
@@ -84,9 +84,10 @@
     (b'\x00\x00\x00\x00\x00\x00\x00\x01', 'TimeType', Time(1))
 )
 
-ordered_map_value = OrderedMap([(u'\u307fbob', 199),
-                                (u'', -1),
-                                (u'\\', 0)])
+ordered_map_value = OrderedMapSerializedKey(UTF8Type, 2)
+ordered_map_value._insert(u'\u307fbob', 199)
+ordered_map_value._insert(u'', -1)
+ordered_map_value._insert(u'\\', 0)
 
 # these following entries work for me right now, but they're dependent on
 # vagaries of internal python ordering for unordered types