Merge "Fix when we set state in Spliterator"

swift/common/swob.py

@@ -1248,6 +1248,8 @@ class Response(object):
             if ranges == []:
                 self.status = 416
                 close_if_possible(app_iter)
+                # Setting body + app_iter to None makes us emit the default
+                # body text from RESPONSE_REASONS.
                 body = None
                 app_iter = None
             elif ranges:

swift/common/utils.py

@@ -3428,8 +3428,9 @@ class Spliterator(object):
 
     def take(self, n):
         if self._iterator_in_progress:
-            raise ValueError("cannot call take() again until the first"
+            raise ValueError(
-                             " iterator is exhausted")
+                "cannot call take() again until the first iterator is"
+                " exhausted (has raised StopIteration)")
         self._iterator_in_progress = True
 
         try:
@@ -3451,14 +3452,16 @@ class Spliterator(object):
                 llen = len(self.leftovers) - self.leftovers_index
                 if llen <= n:
                     n -= llen
-                    yield self.leftovers[self.leftovers_index:]
+                    to_yield = self.leftovers[self.leftovers_index:]
                     self.leftovers = None
                     self.leftovers_index = 0
+                    yield to_yield
                 else:
-                    yield self.leftovers[
+                    to_yield = self.leftovers[
                         self.leftovers_index:(self.leftovers_index + n)]
                     self.leftovers_index += n
                     n = 0
+                    yield to_yield
 
             while n > 0:
                 chunk = next(self.input_iterator)
@@ -3467,9 +3470,9 @@ class Spliterator(object):
                     n -= cl
                     yield chunk
                 else:
-                    yield chunk[:n]
                     self.leftovers = chunk
                     self.leftovers_index = n
+                    yield chunk[:n]
                     n = 0
         finally:
             self._iterator_in_progress = False

test/unit/common/test_utils.py

@@ -5566,6 +5566,36 @@ class TestSpliterator(unittest.TestCase):
         t2 = si.take(20)
         self.assertRaises(ValueError, next, t2)
 
+    def test_closing(self):
+        input_chunks = ["abcd", "efg", "hij"]
+
+        si = utils.Spliterator(input_chunks)
+        it = si.take(3)  # shorter than first chunk
+        self.assertEqual(next(it), 'abc')
+        it.close()
+        self.assertEqual(list(si.take(20)), ['d', 'efg', 'hij'])
+
+        si = utils.Spliterator(input_chunks)
+        self.assertEqual(list(si.take(1)), ['a'])
+        it = si.take(1)  # still shorter than first chunk
+        self.assertEqual(next(it), 'b')
+        it.close()
+        self.assertEqual(list(si.take(20)), ['cd', 'efg', 'hij'])
+
+        si = utils.Spliterator(input_chunks)
+        it = si.take(6)  # longer than first chunk, shorter than first + second
+        self.assertEqual(next(it), 'abcd')
+        self.assertEqual(next(it), 'ef')
+        it.close()
+        self.assertEqual(list(si.take(20)), ['g', 'hij'])
+
+        si = utils.Spliterator(input_chunks)
+        self.assertEqual(list(si.take(2)), ['ab'])
+        it = si.take(3)  # longer than rest of chunk
+        self.assertEqual(next(it), 'cd')
+        it.close()
+        self.assertEqual(list(si.take(20)), ['efg', 'hij'])
+
 
 class TestParseContentRange(unittest.TestCase):
     def test_good(self):