Fix when we set state in Spliterator

Also clean up a comment and some exception text

Change-Id: I1e7755cc0468f9a3ba96a0dd24868f09a10c3df0
Related-Change: I24716e3271cf3370642e3755447e717fd7d9957c
This commit is contained in:
Tim Burke 2016-09-19 17:37:47 -07:00
parent 4bcd3d7f6d
commit 2ec4189e37
3 changed files with 40 additions and 5 deletions

View File

@ -1246,6 +1246,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:

View File

@ -3319,8 +3319,9 @@ class Spliterator(object):
def take(self, n):
if self._iterator_in_progress:
raise ValueError("cannot call take() again until the first"
" iterator is exhausted")
raise ValueError(
"cannot call take() again until the first iterator is"
" exhausted (has raised StopIteration)")
self._iterator_in_progress = True
try:
@ -3342,14 +3343,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)
@ -3358,9 +3361,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

View File

@ -5396,6 +5396,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):