DictCursor and RealDictCursor rows maintain columns order

Close #177.

3 files changed, 95 insertions, 6 deletions

diff --git a/NEWS b/NEWS
index c3268d4..3dd6cd9 100644
--- a/NEWS
+++ b/NEWS
@@ -7,6 +7,8 @@ What's new in psycopg 2.8
 New features:
 
 - Added `~psycopg2.extensions.encrypt_password()` function (:ticket:`#576`).
+- `~psycopg2.extras.DictCursor` and `~psycopg2.extras.RealDictCursor` rows
+  maintain columns order (:ticket:`#177`).
 
 Other changes:
 
diff --git a/lib/extras.py b/lib/extras.py
index 904c5f0..ff32ab6 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -29,7 +29,7 @@ import os as _os
 import sys as _sys
 import time as _time
 import re as _re
-from collections import namedtuple
+from collections import namedtuple, OrderedDict
 
 try:
     import logging as _logging
@@ -140,12 +140,12 @@ class DictCursor(DictCursorBase):
         self._prefetch = 1
 
     def execute(self, query, vars=None):
-        self.index = {}
+        self.index = OrderedDict()
         self._query_executed = 1
         return super(DictCursor, self).execute(query, vars)
 
     def callproc(self, procname, vars=None):
-        self.index = {}
+        self.index = OrderedDict()
         self._query_executed = 1
         return super(DictCursor, self).callproc(procname, vars)
 
@@ -193,7 +193,7 @@ class DictRow(list):
             return default
 
     def copy(self):
-        return dict(self.items())
+        return OrderedDict(self.items())
 
     def __contains__(self, x):
         return x in self._index
@@ -282,6 +282,32 @@ class RealDictRow(dict):
         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)
+
+    if _sys.version_info[0] < 3:
+        iterkeys = keys
+        itervalues = values
+        iteritems = items
+
+        def keys(self):
+            return list(self.iterkeys())
+
+        def values(self):
+            return list(self.itervalues())
+
+        def items(self):
+            return list(self.iteritems())
+
 
 class NamedTupleConnection(_connection):
     """A connection that uses `NamedTupleCursor` automatically."""
diff --git a/tests/test_extras_dictcursor.py b/tests/test_extras_dictcursor.py
index a9201f1..1d5dfd0 100755
--- a/tests/test_extras_dictcursor.py
+++ b/tests/test_extras_dictcursor.py
@@ -15,6 +15,7 @@
 # License for more details.
 
 import time
+import pickle
 from datetime import timedelta
 import psycopg2
 import psycopg2.extras
@@ -140,7 +141,6 @@ class ExtrasDictCursorTests(_DictCursorBase):
         self.failUnless(row[0] == 'bar')
 
     def testPickleDictRow(self):
-        import pickle
         curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
         curs.execute("select 10 as a, 20 as b")
         r = curs.fetchone()
@@ -184,6 +184,37 @@ class ExtrasDictCursorTests(_DictCursorBase):
         self.assert_(not isinstance(r.items(), list))
         self.assertEqual(len(list(r.items())), 2)
 
+    def test_order(self):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+        curs.execute("select 5 as foo, 4 as bar, 33 as baz, 2 as qux")
+        r = curs.fetchone()
+        self.assertEqual(list(r), [5, 4, 33, 2])
+        self.assertEqual(list(r.keys()), ['foo', 'bar', 'baz', 'qux'])
+        self.assertEqual(list(r.values()), [5, 4, 33, 2])
+        self.assertEqual(list(r.items()),
+            [('foo', 5), ('bar', 4), ('baz', 33), ('qux', 2)])
+
+        r1 = pickle.loads(pickle.dumps(r))
+        self.assertEqual(list(r1), list(r))
+        self.assertEqual(list(r1.keys()), list(r.keys()))
+        self.assertEqual(list(r1.values()), list(r.values()))
+        self.assertEqual(list(r1.items()), list(r.items()))
+
+    @skip_from_python(3)
+    def test_order_iter(self):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+        curs.execute("select 5 as foo, 4 as bar, 33 as baz, 2 as qux")
+        r = curs.fetchone()
+        self.assertEqual(list(r.iterkeys()), ['foo', 'bar', 'baz', 'qux'])
+        self.assertEqual(list(r.itervalues()), [5, 4, 33, 2])
+        self.assertEqual(list(r.iteritems()),
+            [('foo', 5), ('bar', 4), ('baz', 33), ('qux', 2)])
+
+        r1 = pickle.loads(pickle.dumps(r))
+        self.assertEqual(list(r1.iterkeys()), list(r.iterkeys()))
+        self.assertEqual(list(r1.itervalues()), list(r.itervalues()))
+        self.assertEqual(list(r1.iteritems()), list(r.iteritems()))
+
 
 class ExtrasDictCursorRealTests(_DictCursorBase):
     def testDictCursorWithPlainCursorRealFetchOne(self):
@@ -216,7 +247,6 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
         self.failUnless(row['foo'] == 'bar')
 
     def testPickleRealDictRow(self):
-        import pickle
         curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
         curs.execute("select 10 as a, 20 as b")
         r = curs.fetchone()
@@ -293,6 +323,37 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
         self.assert_(not isinstance(r.items(), list))
         self.assertEqual(len(list(r.items())), 2)
 
+    def test_order(self):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
+        curs.execute("select 5 as foo, 4 as bar, 33 as baz, 2 as qux")
+        r = curs.fetchone()
+        self.assertEqual(list(r), ['foo', 'bar', 'baz', 'qux'])
+        self.assertEqual(list(r.keys()), ['foo', 'bar', 'baz', 'qux'])
+        self.assertEqual(list(r.values()), [5, 4, 33, 2])
+        self.assertEqual(list(r.items()),
+            [('foo', 5), ('bar', 4), ('baz', 33), ('qux', 2)])
+
+        r1 = pickle.loads(pickle.dumps(r))
+        self.assertEqual(list(r1), list(r))
+        self.assertEqual(list(r1.keys()), list(r.keys()))
+        self.assertEqual(list(r1.values()), list(r.values()))
+        self.assertEqual(list(r1.items()), list(r.items()))
+
+    @skip_from_python(3)
+    def test_order_iter(self):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
+        curs.execute("select 5 as foo, 4 as bar, 33 as baz, 2 as qux")
+        r = curs.fetchone()
+        self.assertEqual(list(r.iterkeys()), ['foo', 'bar', 'baz', 'qux'])
+        self.assertEqual(list(r.itervalues()), [5, 4, 33, 2])
+        self.assertEqual(list(r.iteritems()),
+            [('foo', 5), ('bar', 4), ('baz', 33), ('qux', 2)])
+
+        r1 = pickle.loads(pickle.dumps(r))
+        self.assertEqual(list(r1.iterkeys()), list(r.iterkeys()))
+        self.assertEqual(list(r1.itervalues()), list(r.itervalues()))
+        self.assertEqual(list(r1.iteritems()), list(r.iteritems()))
+
 
 class NamedTupleCursorTest(ConnectingTestCase):
     def setUp(self):