refactor: OrderedDict copy implement

Author: fatelei

Lib/test/test_ordered_dict.py

@@ -730,32 +730,6 @@ def test_merge_operator(self):
         with self.assertRaises(ValueError):
             a |= "BAD"
 
-    def test_getitem_re_entrant_clear_during_copy(self):
-        class Evil(self.OrderedDict):
-            def __getitem__(self, key):
-                super().clear()
-                return None
-
-        evil_dict = Evil([(i, i) for i in range(4)])
-        result = evil_dict.copy()
-
-        self.assertEqual(len(result), 4)
-
-    def test_getitem_re_entrant_modify_during_copy(self):
-        class Modifier(self.OrderedDict):
-            def __getitem__(self, key):
-                self['new_key'] = 'new_value'
-                return super().__getitem__(key)
-
-        original = Modifier([(1, 'one'), (2, 'two')])
-        result = original.copy()
-
-        self.assertIn(1, result)
-        self.assertIn(2, result)
-        self.assertEqual(result[1], 'one')
-        self.assertEqual(result[2], 'two')
-        self.assertEqual(result["new_key"], "new_value")
-
     @support.cpython_only
     def test_ordered_dict_items_result_gc(self):
         # bpo-42536: OrderedDict.items's tuple-reuse speed trick breaks the GC's
@@ -900,6 +874,26 @@ def side_effect(self):
         self.assertDictEqual(dict1, dict.fromkeys((0, 4.2)))
         self.assertDictEqual(dict2, dict.fromkeys((0, Key(), 4.2)))
 
+    def test_getitem_re_entrant_clear_during_copy(self):
+        class Evil(self.OrderedDict):
+            def __getitem__(self, key):
+                super().clear()
+                return None
+
+        evil_dict = Evil([(i, i) for i in range(4)])
+        with self.assertRaises(RuntimeError):
+            result = evil_dict.copy()
+
+    def test_getitem_re_entrant_modify_during_copy(self):
+        class Modifier(self.OrderedDict):
+            def __getitem__(self, key):
+                self['new_key'] = 'new_value'
+                return super().__getitem__(key)
+
+        original = Modifier([(1, 'one'), (2, 'two')])
+        with self.assertRaises(RuntimeError):
+            result = original.copy()
+
 
 @unittest.skipUnless(c_coll, 'requires the C version of the collections module')
 class CPythonOrderedDictTests(OrderedDictTests,

Objects/odictobject.c

@@ -1266,51 +1266,39 @@ OrderedDict_copy_impl(PyObject *od)
         }
     }
     else {
-        Py_ssize_t i, size = PyODict_Size((PyODictObject *)od);
-        PyObject **keys;
-        if (size < 0) {
-            goto fail;
-        }
-
-        if (size == 0) {
-            return od_copy;
-        }
-
-        keys = PyMem_Malloc(size * sizeof(PyObject *));
-        if (keys == NULL) {
-            PyErr_NoMemory();
-            goto fail;
-        }
+        PyODictObject *self = _PyODictObject_CAST(od);
+        size_t state = self->od_state;
+        _ODictNode *cur = _odict_FIRST(od);
 
-        i = 0;
-        _odict_FOREACH(od, node) {
-            keys[i] = _odictnode_KEY(node);
-            Py_INCREF(keys[i]);
-            i++;
-        }
-
-        for (i = 0; i < size; i++) {
-            int res;
-            PyObject *value = PyObject_GetItem((PyObject *)od, keys[i]);
+        while (cur != NULL) {
+            if (self->od_state != state) {
+                PyErr_SetString(PyExc_RuntimeError,
+                                "OrderedDict mutated during iteration");
+                goto fail;
+            }
+            PyObject *key = Py_NewRef(_odictnode_KEY(cur));
+            PyObject *value = PyObject_GetItem((PyObject *)od, key);
             if (value == NULL) {
-                for (Py_ssize_t j = 0; j < size; j++) {
-                    Py_DECREF(keys[j]);
-                }
-                PyMem_Free(keys);
+                Py_DECREF(key);
                 goto fail;
             }
-            res = PyObject_SetItem((PyObject *)od_copy, keys[i], value);
-            Py_DECREF(value);
-            Py_DECREF(keys[i]);
-            if (res != 0) {
-                for (; i < size; i++) {
-                    Py_DECREF(keys[i]);
-                }
-                PyMem_Free(keys);
+            if (self->od_state != state) {
+                Py_DECREF(key);
+                Py_DECREF(value);
+                PyErr_SetString(PyExc_RuntimeError,
+                                "OrderedDict mutated during iteration");
                 goto fail;
             }
+            if (PyObject_SetItem((PyObject *)od_copy, key, value) != 0) {
+                Py_DECREF(key);
+                Py_DECREF(value);
+                goto fail;
+            }
+            Py_DECREF(key);
+            Py_DECREF(value);
+
+            cur = _odictnode_NEXT(cur);
         }
-        PyMem_Free(keys);
     }
     return od_copy;