From cc815e8e8d33d5f39e602eebd6f7341952dfe058 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 6 Apr 2019 19:39:42 +0100
Subject: [PATCH] RealDictRow inherits from OrderedDict

Now its state is unmodified, so apart from special-casing creation
and initial population can work unmodified, and all the desired
properties just work (modifiability, picklability...)

Close #886.
---
 NEWS                            |  2 +-
 lib/extras.py                   | 76 +++++++++++----------------------
 tests/test_extras_dictcursor.py | 17 +++++++-
 3 files changed, 43 insertions(+), 52 deletions(-)

diff --git a/NEWS b/NEWS
index 7aa146e5..c2500fda 100644
--- a/NEWS
+++ b/NEWS
@@ -4,7 +4,7 @@ Current release
 What's new in psycopg 2.8.1
 ---------------------------
 
-- Fixed `RealDictRow.pop()` (:ticket:`#886`).
+- Fixed `RealDictRow` modifiability (:ticket:`#886`).
 - Fixed "there's no async cursor" error polling a connection with no cursor
   (:ticket:`#887`).
 
diff --git a/lib/extras.py b/lib/extras.py
index c5294cef..42c0d3c9 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, *args, **kwargs):
+        if args and isinstance(args[0], _cursor):
+            cursor = args[0]
+            args = args[1:]
+        else:
+            cursor = None
 
-    def __init__(self, cursor):
-        super(RealDictRow, self).__init__()
-        # Required for named cursors
-        if cursor.description and not cursor.column_mapping:
-            cursor._build_index()
+        super(RealDictRow, self).__init__(*args, **kwargs)
 
-        self._column_mapping = cursor.column_mapping
+        if cursor is not None:
+            # Required for named cursors
+            if cursor.description and not cursor.column_mapping:
+                cursor._build_index()
 
-    def __setitem__(self, name, value):
-        if type(name) == int:
-            name = self._column_mapping[name]
-        super(RealDictRow, self).__setitem__(name, value)
+            # 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 __getstate__(self):
-        return self.copy(), self._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 __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 items(self):
-            return list(self.iteritems())
+        super(RealDictRow, self).__setitem__(key, value)
 
 
 class NamedTupleConnection(_connection):
diff --git a/tests/test_extras_dictcursor.py b/tests/test_extras_dictcursor.py
index 119110e8..c7f09d54 100755
--- a/tests/test_extras_dictcursor.py
+++ b/tests/test_extras_dictcursor.py
@@ -265,7 +265,6 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
         self.assertEqual(r, r1)
         self.assertEqual(r['a'], r1['a'])
         self.assertEqual(r['b'], r1['b'])
-        self.assertEqual(r._column_mapping, r1._column_mapping)
 
     def testDictCursorRealWithNamedCursorFetchOne(self):
         self._testWithNamedCursorReal(lambda curs: curs.fetchone())
@@ -377,6 +376,22 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
         self.assertEqual(r.pop('b', None), None)
         self.assertRaises(KeyError, r.pop, 'b')
 
+    def test_mod(self):
+        curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
+        curs.execute("select 1 as a, 2 as b, 3 as c")
+        r = curs.fetchone()
+        r['d'] = 4
+        self.assertEqual(list(r), ['a', 'b', 'c', 'd'])
+        self.assertEqual(list(r.keys()), ['a', 'b', 'c', 'd'])
+        self.assertEqual(list(r.values()), [1, 2, 3, 4])
+        self.assertEqual(list(
+            r.items()), [('a', 1), ('b', 2), ('c', 3), ('d', 4)])
+
+        assert r['a'] == 1
+        assert r['b'] == 2
+        assert r['c'] == 3
+        assert r['d'] == 4
+
 
 class NamedTupleCursorTest(ConnectingTestCase):
     def setUp(self):