Guard iterations (#100)

* Fix #99: guard iterations

* Update CHANGES

Author: asvetlov

CHANGES.rst

@@ -1,3 +1,9 @@
+3.1.0 (2017-xx-xx)
+------------------
+
+* Fix #99: raise `RuntimeError` on dict iterations if the dict was changed
+
+
 3.0.0 (2017-06-21)
 ------------------
 

multidict/_multidict.pyx

@@ -638,16 +638,20 @@ cdef class _ItemsIter:
     cdef _Impl _impl
     cdef int _current
     cdef int _len
+    cdef unsigned long long _version
 
     def __cinit__(self, _Impl impl):
         self._impl = impl
         self._current = 0
-        self._len = len(self._impl._items)
+        self._version = impl._version
+        self._len = len(impl._items)
 
     def __iter__(self):
         return self
 
     def __next__(self):
+        if self._version != self._impl._version:
+            raise RuntimeError("Dictionary changed during iteration")
         if self._current == self._len:
             raise StopIteration
         item = <_Pair>self._impl._items[self._current]
@@ -700,16 +704,20 @@ cdef class _ValuesIter:
     cdef _Impl _impl
     cdef int _current
     cdef int _len
+    cdef unsigned long long _version
 
     def __cinit__(self, _Impl impl):
         self._impl = impl
         self._current = 0
-        self._len = len(self._impl._items)
+        self._len = len(impl._items)
+        self._version = impl._version
 
     def __iter__(self):
         return self
 
     def __next__(self):
+        if self._version != self._impl._version:
+            raise RuntimeError("Dictionary changed during iteration")
         if self._current == self._len:
             raise StopIteration
         item = <_Pair>self._impl._items[self._current]
@@ -747,16 +755,20 @@ cdef class _KeysIter:
     cdef _Impl _impl
     cdef int _current
     cdef int _len
+    cdef unsigned long long _version
 
     def __cinit__(self, _Impl impl):
         self._impl = impl
         self._current = 0
         self._len = len(self._impl._items)
+        self._version = impl._version
 
     def __iter__(self):
         return self
 
     def __next__(self):
+        if self._version != self._impl._version:
+            raise RuntimeError("Dictionary changed during iteration")
         if self._current == self._len:
             raise StopIteration
         item = <_Pair>self._impl._items[self._current]

multidict/_multidict_py.py

@@ -375,6 +375,7 @@ class _ViewBase:
 
     def __init__(self, impl):
         self._impl = impl
+        self._version = impl._version
 
     def __len__(self):
         return len(self._impl._items)
@@ -392,6 +393,8 @@ def __contains__(self, item):
 
     def __iter__(self):
         for i, k, v in self._impl._items:
+            if self._version != self._impl._version:
+                raise RuntimeError("Dictionary changed during iteration")
             yield k, v
 
     def __repr__(self):
@@ -412,6 +415,8 @@ def __contains__(self, value):
 
     def __iter__(self):
         for item in self._impl._items:
+            if self._version != self._impl._version:
+                raise RuntimeError("Dictionary changed during iteration")
             yield item[2]
 
     def __repr__(self):
@@ -432,6 +437,8 @@ def __contains__(self, key):
 
     def __iter__(self):
         for item in self._impl._items:
+            if self._version != self._impl._version:
+                raise RuntimeError("Dictionary changed during iteration")
             yield item[1]
 
     def __repr__(self):

tests/test_guard.py

@@ -0,0 +1,34 @@
+import pytest
+
+from multidict._multidict import MultiDict
+from multidict._multidict_py import MultiDict as PyMultiDict
+
+
+@pytest.fixture(params=[MultiDict, PyMultiDict],
+                ids=['MultiDict', 'PyMultiDict'])
+def cls(request):
+    return request.param
+
+
+def test_guard_items(cls):
+    md = cls({'a': 'b'})
+    it = iter(md.items())
+    md['a'] = 'c'
+    with pytest.raises(RuntimeError):
+        next(it)
+
+
+def test_guard_keys(cls):
+    md = cls({'a': 'b'})
+    it = iter(md.keys())
+    md['a'] = 'c'
+    with pytest.raises(RuntimeError):
+        next(it)
+
+
+def test_guard_values(cls):
+    md = cls({'a': 'b'})
+    it = iter(md.values())
+    md['a'] = 'c'
+    with pytest.raises(RuntimeError):
+        next(it)