Fixed pickling of DictRow objects too

3 files changed, 23 insertions, 2 deletions

diff --git a/NEWS b/NEWS
index 725a1ec..e49f384 100644
--- a/NEWS
+++ b/NEWS
@@ -37,7 +37,7 @@ What's new in psycopg 2.4.6
   - Fixed pickling of FixedOffsetTimezone objects (ticket #135).
   - Release the GIL around PQgetResult calls after COPY (ticket #140).
   - Fixed empty strings handling in composite caster (ticket #141).
-  - Fixed pickling of RealDictRow objects.
+  - Fixed pickling of DictRow and RealDictRow objects.
 
 
 What's new in psycopg 2.4.5
diff --git a/lib/extras.py b/lib/extras.py
index 0c4df56..1924908 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -184,7 +184,14 @@ class DictRow(list):
     def __contains__(self, x):
         return x in self._index
 
-    # grop the crusty Py2 methods
+    def __getstate__(self):
+        return self[:], self._index.copy()
+
+    def __setstate__(self, data):
+        self[:] = data[0]
+        self._index = data[1]
+
+    # drop the crusty Py2 methods
     if sys.version_info[0] > 2:
         items = iteritems; del iteritems
         keys = iterkeys; del iterkeys
diff --git a/tests/test_extras_dictcursor.py b/tests/test_extras_dictcursor.py
index e365c10..4846dea 100755
--- a/tests/test_extras_dictcursor.py
+++ b/tests/test_extras_dictcursor.py
@@ -205,6 +205,20 @@ class ExtrasDictCursorTests(unittest.TestCase):
         for i, r in enumerate(curs):
             self.assertEqual(i + 1, curs.rownumber)
 
+    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()
+        d = pickle.dumps(r)
+        r1 = pickle.loads(d)
+        self.assertEqual(r, r1)
+        self.assertEqual(r[0], r1[0])
+        self.assertEqual(r[1], r1[1])
+        self.assertEqual(r['a'], r1['a'])
+        self.assertEqual(r['b'], r1['b'])
+        self.assertEqual(r._index, r1._index)
+
     def testPickleRealDictRow(self):
         import pickle
         curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)