|
| 1 | +import _pyio |
1 | 2 | import array
|
2 | 3 | import contextlib
|
3 | 4 | import importlib.util
|
@@ -3491,5 +3492,87 @@ def test_too_short(self):
|
3491 | 3492 | b"zzz", zipfile._Extra.strip(b"zzz", (self.ZIP64_EXTRA,)))
|
3492 | 3493 |
|
3493 | 3494 |
|
| 3495 | +class StatIO(_pyio.BytesIO): |
| 3496 | + """Buffer which remembers the number of bytes that were read.""" |
| 3497 | + |
| 3498 | + def __init__(self): |
| 3499 | + super().__init__() |
| 3500 | + self.bytes_read = 0 |
| 3501 | + |
| 3502 | + def read(self, size=-1): |
| 3503 | + bs = super().read(size) |
| 3504 | + self.bytes_read += len(bs) |
| 3505 | + return bs |
| 3506 | + |
| 3507 | + |
| 3508 | +class StoredZipExtFileRandomReadTest(unittest.TestCase): |
| 3509 | + """Tests whether an uncompressed, unencrypted zip entry can be randomly |
| 3510 | + seek and read without reading redundant bytes.""" |
| 3511 | + def test_stored_seek_and_read(self): |
| 3512 | + |
| 3513 | + sio = StatIO() |
| 3514 | + # 20000 bytes |
| 3515 | + txt = b'0123456789' * 2000 |
| 3516 | + |
| 3517 | + # The seek length must be greater than ZipExtFile.MIN_READ_SIZE |
| 3518 | + # as `ZipExtFile._read2()` reads in blocks of this size and we |
| 3519 | + # need to seek out of the buffered data |
| 3520 | + read_buffer_size = zipfile.ZipExtFile.MIN_READ_SIZE |
| 3521 | + self.assertGreaterEqual(10002, read_buffer_size) # for forward seek test |
| 3522 | + self.assertGreaterEqual(5003, read_buffer_size) # for backward seek test |
| 3523 | + # The read length must be less than MIN_READ_SIZE, since we assume that |
| 3524 | + # only 1 block is read in the test. |
| 3525 | + read_length = 100 |
| 3526 | + self.assertGreaterEqual(read_buffer_size, read_length) # for read() calls |
| 3527 | + |
| 3528 | + with zipfile.ZipFile(sio, "w", compression=zipfile.ZIP_STORED) as zipf: |
| 3529 | + zipf.writestr("foo.txt", txt) |
| 3530 | + |
| 3531 | + # check random seek and read on a file |
| 3532 | + with zipfile.ZipFile(sio, "r") as zipf: |
| 3533 | + with zipf.open("foo.txt", "r") as fp: |
| 3534 | + # Test this optimized read hasn't rewound and read from the |
| 3535 | + # start of the file (as in the case of the unoptimized path) |
| 3536 | + |
| 3537 | + # forward seek |
| 3538 | + old_count = sio.bytes_read |
| 3539 | + forward_seek_len = 10002 |
| 3540 | + current_pos = 0 |
| 3541 | + fp.seek(forward_seek_len, os.SEEK_CUR) |
| 3542 | + current_pos += forward_seek_len |
| 3543 | + self.assertEqual(fp.tell(), current_pos) |
| 3544 | + self.assertEqual(fp._left, fp._compress_left) |
| 3545 | + arr = fp.read(read_length) |
| 3546 | + current_pos += read_length |
| 3547 | + self.assertEqual(fp.tell(), current_pos) |
| 3548 | + self.assertEqual(arr, txt[current_pos - read_length:current_pos]) |
| 3549 | + self.assertEqual(fp._left, fp._compress_left) |
| 3550 | + read_count = sio.bytes_read - old_count |
| 3551 | + self.assertLessEqual(read_count, read_buffer_size) |
| 3552 | + |
| 3553 | + # backward seek |
| 3554 | + old_count = sio.bytes_read |
| 3555 | + backward_seek_len = 5003 |
| 3556 | + fp.seek(-backward_seek_len, os.SEEK_CUR) |
| 3557 | + current_pos -= backward_seek_len |
| 3558 | + self.assertEqual(fp.tell(), current_pos) |
| 3559 | + self.assertEqual(fp._left, fp._compress_left) |
| 3560 | + arr = fp.read(read_length) |
| 3561 | + current_pos += read_length |
| 3562 | + self.assertEqual(fp.tell(), current_pos) |
| 3563 | + self.assertEqual(arr, txt[current_pos - read_length:current_pos]) |
| 3564 | + self.assertEqual(fp._left, fp._compress_left) |
| 3565 | + read_count = sio.bytes_read - old_count |
| 3566 | + self.assertLessEqual(read_count, read_buffer_size) |
| 3567 | + |
| 3568 | + # eof flags test |
| 3569 | + fp.seek(0, os.SEEK_END) |
| 3570 | + fp.seek(12345, os.SEEK_SET) |
| 3571 | + current_pos = 12345 |
| 3572 | + arr = fp.read(read_length) |
| 3573 | + current_pos += read_length |
| 3574 | + self.assertEqual(arr, txt[current_pos - read_length:current_pos]) |
| 3575 | + |
| 3576 | + |
3494 | 3577 | if __name__ == "__main__":
|
3495 | 3578 | unittest.main()
|
0 commit comments