diff options
Diffstat (limited to 'lib/extras.py')
-rw-r--r-- | lib/extras.py | 84 |
1 files changed, 30 insertions, 54 deletions
diff --git a/lib/extras.py b/lib/extras.py index c5294ce..42c0d3c 100644 --- a/lib/extras.py +++ b/lib/extras.py @@ -253,63 +253,39 @@ class RealDictCursor(DictCursorBase): self._query_executed = False -class RealDictRow(dict): +class RealDictRow(OrderedDict): """A `!dict` subclass representing a data record.""" - __slots__ = ('_column_mapping',) - - def __init__(self, cursor): - super(RealDictRow, self).__init__() - # Required for named cursors - if cursor.description and not cursor.column_mapping: - cursor._build_index() - - self._column_mapping = cursor.column_mapping - - def __setitem__(self, name, value): - if type(name) == int: - name = self._column_mapping[name] - super(RealDictRow, self).__setitem__(name, value) - - def __getstate__(self): - return self.copy(), self._column_mapping[:] - - def __setstate__(self, data): - self.update(data[0]) - self._column_mapping = data[1] - - def __iter__(self): - return iter(self._column_mapping) - - def keys(self): - return iter(self._column_mapping) - - def values(self): - return (self[k] for k in self._column_mapping) - - def items(self): - return ((k, self[k]) for k in self._column_mapping) - - def pop(self, key, *args): - found = key in self - rv = super(RealDictRow, self).pop(key, *args) - if found: - self._column_mapping.remove(key) - return rv - - if PY2: - iterkeys = keys - itervalues = values - iteritems = items - - def keys(self): - return list(self.iterkeys()) - - def values(self): - return list(self.itervalues()) + def __init__(self, *args, **kwargs): + if args and isinstance(args[0], _cursor): + cursor = args[0] + args = args[1:] + else: + cursor = None + + super(RealDictRow, self).__init__(*args, **kwargs) + + if cursor is not None: + # Required for named cursors + if cursor.description and not cursor.column_mapping: + cursor._build_index() + + # Store the cols mapping in the dict itself until the row is fully + # populated, so we don't need to add attributes to the class + # (hence keeping its maintenance, special pickle support, etc.) + self[RealDictRow] = cursor.column_mapping + + def __setitem__(self, key, value): + if RealDictRow in self: + # We are in the row building phase + mapping = self[RealDictRow] + super(RealDictRow, self).__setitem__(mapping[key], value) + if len(self) == len(mapping) + 1: + # Row building finished + del self[RealDictRow] + return - def items(self): - return list(self.iteritems()) + super(RealDictRow, self).__setitem__(key, value) class NamedTupleConnection(_connection): |