Restore multidict internal state on exception during updating (#1215)

Author: asvetlov

CHANGES/1215.bugfix.rst

@@ -0,0 +1,2 @@
+If :meth:`multidict.MultiDict.extend`, :meth:`multidict.MultiDict.merge`, or :meth:`multidict.MultiDict.update` raises an exception, now the multidict internal state is correctly restored.
+Patch by :user:`asvetlov`.

multidict/_multidict.c

@@ -118,15 +118,18 @@ _multidict_extend(MultiDictObject *self, PyObject *arg, PyObject *kwds,
     }
 
     if (op != Extend) {  // Update or Merge
-        if (md_post_update(self) < 0) {
-            goto fail;
-        }
+        md_post_update(self);
     }
 
     ASSERT_CONSISTENT(self, false);
     Py_CLEAR(seq);
     return 0;
 fail:
+    if (op != Extend) {  // Update or Merge
+        // Cleanup soft-deleted items
+        md_post_update(self);
+    }
+    ASSERT_CONSISTENT(self, false);
     Py_CLEAR(seq);
     return -1;
 }

multidict/_multidict_py.py

@@ -445,11 +445,11 @@ def _identity(self, key: str) -> str:
         if isinstance(key, istr):
             ret = key.__istr_identity__
             if ret is None:
-                ret = key.title()
+                ret = key.lower()
                 key.__istr_identity__ = ret
             return ret
         if isinstance(key, str):
-            return key.title()
+            return key.lower()
         else:
             raise TypeError("MultiDict keys should be either str or subclasses of str")
 
@@ -632,13 +632,13 @@ def __init__(self, arg: MDArg[_V] = None, /, **kwargs: _V):
                 self._from_md(md)
                 return
 
-        items = self._parse_args(arg, kwargs)
-        log2_size = estimate_log2_keysize(len(items))
+        it = self._parse_args(arg, kwargs)
+        log2_size = estimate_log2_keysize(cast(int, next(it)))
         if log2_size > 17:  # pragma: no cover
             # Don't overallocate really huge keys space in init
             log2_size = 17
         self._keys: _HtKeys[_V] = _HtKeys.new(log2_size, [])
-        self._extend_items(items)
+        self._extend_items(cast(Iterator[_Entry[_V]], it))
 
     def _from_md(self, md: "MultiDict[_V]") -> None:
         # Copy everything as-is without compacting the new multidict,
@@ -790,58 +790,57 @@ def extend(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None:
 
         This method must be used instead of update.
         """
-        items = self._parse_args(arg, kwargs)
-        newsize = self._used + len(items)
+        it = self._parse_args(arg, kwargs)
+        newsize = self._used + cast(int, next(it))
         self._resize(estimate_log2_keysize(newsize), False)
-        self._extend_items(items)
+        self._extend_items(cast(Iterator[_Entry[_V]], it))
 
     def _parse_args(
         self,
         arg: MDArg[_V],
         kwargs: Mapping[str, _V],
-    ) -> list[_Entry[_V]]:
+    ) -> Iterator[Union[int, _Entry[_V]]]:
         identity_func = self._identity
         if arg:
             if isinstance(arg, MultiDictProxy):
                 arg = arg._md
             if isinstance(arg, MultiDict):
+                yield len(arg) + len(kwargs)
                 if self._ci is not arg._ci:
-                    items = []
                     for e in arg._keys.iter_entries():
                         identity = identity_func(e.key)
-                        items.append(_Entry(hash(identity), identity, e.key, e.value))
+                        yield _Entry(hash(identity), identity, e.key, e.value)
                 else:
-                    items = [
-                        _Entry(e.hash, e.identity, e.key, e.value)
-                        for e in arg._keys.iter_entries()
-                    ]
+                    for e in arg._keys.iter_entries():
+                        yield _Entry(e.hash, e.identity, e.key, e.value)
                 if kwargs:
                     for key, value in kwargs.items():
                         identity = identity_func(key)
-                        items.append(_Entry(hash(identity), identity, key, value))
+                        yield _Entry(hash(identity), identity, key, value)
             else:
                 if hasattr(arg, "keys"):
                     arg = cast(SupportsKeys[_V], arg)
                     arg = [(k, arg[k]) for k in arg.keys()]
                 if kwargs:
                     arg = list(arg)
                     arg.extend(list(kwargs.items()))
-                items = []
+                try:
+                    yield len(arg) + len(kwargs)  # type: ignore[arg-type]
+                except TypeError:
+                    yield 0
                 for pos, item in enumerate(arg):
                     if not len(item) == 2:
                         raise ValueError(
                             f"multidict update sequence element #{pos}"
                             f"has length {len(item)}; 2 is required"
                         )
                     identity = identity_func(item[0])
-                    items.append(_Entry(hash(identity), identity, item[0], item[1]))
+                    yield _Entry(hash(identity), identity, item[0], item[1])
         else:
-            items = []
+            yield len(kwargs)
             for key, value in kwargs.items():
                 identity = identity_func(key)
-                items.append(_Entry(hash(identity), identity, key, value))
-
-        return items
+                yield _Entry(hash(identity), identity, key, value)
 
     def _extend_items(self, items: Iterable[_Entry[_V]]) -> None:
         for e in items:
@@ -989,19 +988,21 @@ def popitem(self) -> tuple[str, _V]:
 
     def update(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None:
         """Update the dictionary, overwriting existing keys."""
-        items = self._parse_args(arg, kwargs)
-        newsize = self._used + len(items)
+        it = self._parse_args(arg, kwargs)
+        newsize = self._used + cast(int, next(it))
         log2_size = estimate_log2_keysize(newsize)
         if log2_size > 17:  # pragma: no cover
             # Don't overallocate really huge keys space in update,
             # duplicate keys could reduce the resulting anount of entries
             log2_size = 17
         if log2_size > self._keys.log2_size:
             self._resize(log2_size, False)
-        self._update_items(items)
-        self._post_update()
+        try:
+            self._update_items(cast(Iterator[_Entry[_V]], it))
+        finally:
+            self._post_update()
 
-    def _update_items(self, items: list[_Entry[_V]]) -> None:
+    def _update_items(self, items: Iterator[_Entry[_V]]) -> None:
         for entry in items:
             found = False
             hash_ = entry.hash
@@ -1038,19 +1039,21 @@ def _post_update(self) -> None:
 
     def merge(self, arg: MDArg[_V] = None, /, **kwargs: _V) -> None:
         """Merge into the dictionary, adding non-existing keys."""
-        items = self._parse_args(arg, kwargs)
-        newsize = self._used + len(items)
+        it = self._parse_args(arg, kwargs)
+        newsize = self._used + cast(int, next(it))
         log2_size = estimate_log2_keysize(newsize)
         if log2_size > 17:  # pragma: no cover
             # Don't overallocate really huge keys space in update,
             # duplicate keys could reduce the resulting anount of entries
             log2_size = 17
         if log2_size > self._keys.log2_size:
             self._resize(log2_size, False)
-        self._merge_items(items)
-        self._post_update()
+        try:
+            self._merge_items(cast(Iterator[_Entry[_V]], it))
+        finally:
+            self._post_update()
 
-    def _merge_items(self, items: list[_Entry[_V]]) -> None:
+    def _merge_items(self, items: Iterator[_Entry[_V]]) -> None:
         for entry in items:
             hash_ = entry.hash
             identity = entry.identity

multidict/_multilib/hashtable.h

@@ -151,10 +151,10 @@ _ci_key_to_identity(mod_state *state, PyObject *key)
         }
         return ret;
     }
-fail:
     PyErr_SetString(PyExc_TypeError,
                     "CIMultiDict keys should be either str "
                     "or subclasses of str");
+fail:
     return NULL;
 }
 
@@ -1287,7 +1287,7 @@ _md_merge(MultiDictObject *md, Py_hash_t hash, PyObject *identity,
     return -1;
 }
 
-static inline int
+static inline void
 md_post_update(MultiDictObject *md)
 {
     htkeys_t *keys = md->keys;
@@ -1306,14 +1306,11 @@ md_post_update(MultiDictObject *md)
             }
             if (entry->hash == -1) {
                 entry->hash = _unicode_hash(entry->identity);
-                if (entry->hash == -1) {
-                    // hash of string always exists but still
-                    return -1;
-                }
             }
+            assert(entry->hash != -1);
         }
     }
-    return 0;
+    ASSERT_CONSISTENT(md, false);
 }
 
 static inline int

tests/test_mutable_multidict.py

@@ -858,3 +858,17 @@ def test_issue_1195(
         assert md.keys() == md2.keys() - {"User-Agent"}
         md.update([("User-Agent", b"Bacon/1.0")])
         assert md.keys() == md2.keys()
+
+    def test_update_with_crash_in_the_middle(
+        self, case_insensitive_multidict_class: type[CIMultiDict[str]]
+    ) -> None:
+        class Hack(str):
+            def lower(self) -> str:
+                raise RuntimeError
+
+        d = case_insensitive_multidict_class([("a", "a"), ("b", "b")])
+        with pytest.raises(RuntimeError):
+            lst = [("c", "c"), ("a", "a2"), (Hack("b"), "b2")]
+            d.update(lst)
+
+        assert [("a", "a2"), ("b", "b"), ("c", "c")] == list(d.items())