diff --git a/libpathod/rparse.py b/libpathod/rparse.py index 570ee9049..e3136b7b4 100644 --- a/libpathod/rparse.py +++ b/libpathod/rparse.py @@ -36,37 +36,40 @@ def ready_actions(length, lst): return ret +def send_chunk(fp, val, blocksize, start, end): + """ + (start, end): Inclusive lower bound, exclusive upper bound. + """ + for i in range(start, end, blocksize): + fp.write( + val[i:min(i+blocksize, end)] + ) + return end-start + + def write_values(fp, vals, actions, sofar=0, skip=0, blocksize=BLOCKSIZE): """ - vals: A list of values, which may be strings or Value objects. + vals: A list of values, which may be strings or Value objects. actions: A list of (offset, action, arg) tuples. Action may be "pause" or "disconnect". + Both vals and actions are in reverse order, with the first items last. + Return True if connection should disconnect. """ + sofar = 0 try: while vals: - part = vals.pop() - for i in range(skip, len(part), blocksize): - d = part[i:i+blocksize] - if actions and actions[-1][0] < (sofar + len(d)): - p = actions.pop() - offset = p[0]-sofar - vals.append(part) - if p[1] == "pause": - fp.write(d[:offset]) - time.sleep(p[2]) - return write_values( - fp, vals, actions, - sofar=sofar+offset, - skip=i+offset, - blocksize=blocksize - ) - elif p[1] == "disconnect": - fp.write(d[:offset]) - return True - fp.write(d) - sofar += len(d) - skip = 0 + v = vals.pop() + offset = 0 + while actions and actions[-1][0] < (sofar + len(v)): + a = actions.pop() + offset += send_chunk(fp, v, blocksize, offset, a[0]-sofar-offset) + if a[1] == "pause": + time.sleep(a[2]) + elif a[1] == "disconnect": + return True + send_chunk(fp, v, blocksize, offset, len(v)) + sofar += len(v) except tcp.NetLibDisconnect: return True diff --git a/test/test_rparse.py b/test/test_rparse.py index cb1076c28..f4b408b2b 100644 --- a/test/test_rparse.py +++ b/test/test_rparse.py @@ -250,6 +250,18 @@ class TestParseResponse: class TestWriteValues: + def test_send_chunk(self): + v = "foobarfoobar" + for bs in range(1, len(v)+2): + s = cStringIO.StringIO() + rparse.send_chunk(s, v, bs, 0, len(v)) + assert s.getvalue() == v + for start in range(len(v)): + for end in range(len(v)): + s = cStringIO.StringIO() + rparse.send_chunk(s, v, bs, start, end) + assert s.getvalue() == v[start:end] + def test_write_values_disconnects(self): s = cStringIO.StringIO() tst = "foo"*100 @@ -257,11 +269,17 @@ class TestWriteValues: assert not s.getvalue() def test_write_values(self): - tst = "foo"*1025 + tst = "foobarvoing" s = cStringIO.StringIO() rparse.write_values(s, [tst], []) assert s.getvalue() == tst + for bs in range(1, len(tst) + 2): + for off in range(len(tst)): + s = cStringIO.StringIO() + rparse.write_values(s, [tst], [(off, "disconnect")], blocksize=bs) + assert s.getvalue() == tst[:off] + def test_write_values_pauses(self): tst = "".join(str(i) for i in range(10)) for i in range(2, 10):