cqle: add dict access sugar to UserType

This commit is contained in:
Adam Holmberg
2015-03-20 15:33:30 -05:00
parent 8c98eb22d0
commit 79bd4fac17

View File

@@ -1,7 +1,6 @@
import re import re
import six import six
from cassandra.cluster import UserTypeDoesNotExist
from cassandra.util import OrderedDict from cassandra.util import OrderedDict
from cassandra.cqlengine import CQLEngineException from cassandra.cqlengine import CQLEngineException
from cassandra.cqlengine import columns from cassandra.cqlengine import columns
@@ -66,6 +65,43 @@ class BaseUserType(object):
for v in self._values.values(): for v in self._values.values():
v.reset_previous_value() v.reset_previous_value()
def __iter__(self):
for field in self._fields.keys():
yield field
def __getitem__(self, key):
if not isinstance(key, six.string_types):
raise TypeError
if key not in self._fields.keys():
raise KeyError
return getattr(self, key)
def __setitem__(self, key, val):
if not isinstance(key, six.string_types):
raise TypeError
if key not in self._fields.keys():
raise KeyError
return setattr(self, key, val)
def __len__(self):
try:
return self._len
except:
self._len = len(self._columns.keys())
return self._len
def keys(self):
""" Returns a list of column IDs. """
return [k for k in self]
def values(self):
""" Returns list of column values. """
return [self[k] for k in self]
def items(self):
""" Returns a list of column ID/value tuples. """
return [(k, self[k]) for k in self]
@classmethod @classmethod
def register_for_keyspace(cls, keyspace): def register_for_keyspace(cls, keyspace):
connection.register_udt(keyspace, cls.type_name(), cls) connection.register_udt(keyspace, cls.type_name(), cls)