Pārlūkot izejas kodu

Merge pull request #1659 from ThomasWaldmann/fix-iterator

hashindex: fix iterator implementation
enkore 8 gadi atpakaļ
vecāks
revīzija
ab9d761c40
2 mainītis faili ar 14 papildinājumiem un 1 dzēšanām
  1. 10 0
      borg/hashindex.pyx
  2. 4 1
      borg/testsuite/hashindex.py

+ 10 - 0
borg/hashindex.pyx

@@ -168,17 +168,22 @@ cdef class NSKeyIterator:
     cdef HashIndex *index
     cdef HashIndex *index
     cdef const void *key
     cdef const void *key
     cdef int key_size
     cdef int key_size
+    cdef int exhausted
 
 
     def __cinit__(self, key_size):
     def __cinit__(self, key_size):
         self.key = NULL
         self.key = NULL
         self.key_size = key_size
         self.key_size = key_size
+        self.exhausted = 0
 
 
     def __iter__(self):
     def __iter__(self):
         return self
         return self
 
 
     def __next__(self):
     def __next__(self):
+        if self.exhausted:
+            raise StopIteration
         self.key = hashindex_next_key(self.index, <char *>self.key)
         self.key = hashindex_next_key(self.index, <char *>self.key)
         if not self.key:
         if not self.key:
+            self.exhausted = 1
             raise StopIteration
             raise StopIteration
         cdef uint32_t *value = <uint32_t *>(self.key + self.key_size)
         cdef uint32_t *value = <uint32_t *>(self.key + self.key_size)
         cdef uint32_t segment = _le32toh(value[0])
         cdef uint32_t segment = _le32toh(value[0])
@@ -330,17 +335,22 @@ cdef class ChunkKeyIterator:
     cdef HashIndex *index
     cdef HashIndex *index
     cdef const void *key
     cdef const void *key
     cdef int key_size
     cdef int key_size
+    cdef int exhausted
 
 
     def __cinit__(self, key_size):
     def __cinit__(self, key_size):
         self.key = NULL
         self.key = NULL
         self.key_size = key_size
         self.key_size = key_size
+        self.exhausted = 0
 
 
     def __iter__(self):
     def __iter__(self):
         return self
         return self
 
 
     def __next__(self):
     def __next__(self):
+        if self.exhausted:
+            raise StopIteration
         self.key = hashindex_next_key(self.index, <char *>self.key)
         self.key = hashindex_next_key(self.index, <char *>self.key)
         if not self.key:
         if not self.key:
+            self.exhausted = 1
             raise StopIteration
             raise StopIteration
         cdef uint32_t *value = <uint32_t *>(self.key + self.key_size)
         cdef uint32_t *value = <uint32_t *>(self.key + self.key_size)
         cdef uint32_t refcount = _le32toh(value[0])
         cdef uint32_t refcount = _le32toh(value[0])

+ 4 - 1
borg/testsuite/hashindex.py

@@ -83,8 +83,11 @@ class HashIndexTestCase(BaseTestCase):
         idx = NSIndex()
         idx = NSIndex()
         for x in range(100):
         for x in range(100):
             idx[H(x)] = x, x
             idx[H(x)] = x, x
-        all = list(idx.iteritems())
+        iterator = idx.iteritems()
+        all = list(iterator)
         self.assert_equal(len(all), 100)
         self.assert_equal(len(all), 100)
+        # iterator is already exhausted by list():
+        self.assert_raises(StopIteration, next, iterator)
         second_half = list(idx.iteritems(marker=all[49][0]))
         second_half = list(idx.iteritems(marker=all[49][0]))
         self.assert_equal(len(second_half), 50)
         self.assert_equal(len(second_half), 50)
         self.assert_equal(second_half, all[50:])
         self.assert_equal(second_half, all[50:])