|
| 1 | +import _pyio |
1 | 2 | import array
|
2 | 3 | import contextlib
|
3 | 4 | import importlib.util
|
@@ -3454,5 +3455,87 @@ def test_too_short(self):
|
3454 | 3455 | b"zzz", zipfile._Extra.strip(b"zzz", (self.ZIP64_EXTRA,)))
|
3455 | 3456 |
|
3456 | 3457 |
|
| 3458 | +class StatIO(_pyio.BytesIO): |
| 3459 | + """Buffer which remembers the number of bytes that were read.""" |
| 3460 | + |
| 3461 | + def __init__(self): |
| 3462 | + super().__init__() |
| 3463 | + self.bytes_read = 0 |
| 3464 | + |
| 3465 | + def read(self, size=-1): |
| 3466 | + bs = super().read(size) |
| 3467 | + self.bytes_read += len(bs) |
| 3468 | + return bs |
| 3469 | + |
| 3470 | + |
| 3471 | +class StoredZipExtFileRandomReadTest(unittest.TestCase): |
| 3472 | + """Tests whether an uncompressed, unencrypted zip entry can be randomly |
| 3473 | + seek and read without reading redundant bytes.""" |
| 3474 | + def test_stored_seek_and_read(self): |
| 3475 | + |
| 3476 | + sio = StatIO() |
| 3477 | + # 20000 bytes |
| 3478 | + txt = b'0123456789' * 2000 |
| 3479 | + |
| 3480 | + # The seek length must be greater than ZipExtFile.MIN_READ_SIZE |
| 3481 | + # as `ZipExtFile._read2()` reads in blocks of this size and we |
| 3482 | + # need to seek out of the buffered data |
| 3483 | + read_buffer_size = zipfile.ZipExtFile.MIN_READ_SIZE |
| 3484 | + self.assertGreaterEqual(10002, read_buffer_size) # for forward seek test |
| 3485 | + self.assertGreaterEqual(5003, read_buffer_size) # for backward seek test |
| 3486 | + # The read length must be less than MIN_READ_SIZE, since we assume that |
| 3487 | + # only 1 block is read in the test. |
| 3488 | + read_length = 100 |
| 3489 | + self.assertGreaterEqual(read_buffer_size, read_length) # for read() calls |
| 3490 | + |
| 3491 | + with zipfile.ZipFile(sio, "w", compression=zipfile.ZIP_STORED) as zipf: |
| 3492 | + zipf.writestr("foo.txt", txt) |
| 3493 | + |
| 3494 | + # check random seek and read on a file |
| 3495 | + with zipfile.ZipFile(sio, "r") as zipf: |
| 3496 | + with zipf.open("foo.txt", "r") as fp: |
| 3497 | + # Test this optimized read hasn't rewound and read from the |
| 3498 | + # start of the file (as in the case of the unoptimized path) |
| 3499 | + |
| 3500 | + # forward seek |
| 3501 | + old_count = sio.bytes_read |
| 3502 | + forward_seek_len = 10002 |
| 3503 | + current_pos = 0 |
| 3504 | + fp.seek(forward_seek_len, os.SEEK_CUR) |
| 3505 | + current_pos += forward_seek_len |
| 3506 | + self.assertEqual(fp.tell(), current_pos) |
| 3507 | + self.assertEqual(fp._left, fp._compress_left) |
| 3508 | + arr = fp.read(read_length) |
| 3509 | + current_pos += read_length |
| 3510 | + self.assertEqual(fp.tell(), current_pos) |
| 3511 | + self.assertEqual(arr, txt[current_pos - read_length:current_pos]) |
| 3512 | + self.assertEqual(fp._left, fp._compress_left) |
| 3513 | + read_count = sio.bytes_read - old_count |
| 3514 | + self.assertLessEqual(read_count, read_buffer_size) |
| 3515 | + |
| 3516 | + # backward seek |
| 3517 | + old_count = sio.bytes_read |
| 3518 | + backward_seek_len = 5003 |
| 3519 | + fp.seek(-backward_seek_len, os.SEEK_CUR) |
| 3520 | + current_pos -= backward_seek_len |
| 3521 | + self.assertEqual(fp.tell(), current_pos) |
| 3522 | + self.assertEqual(fp._left, fp._compress_left) |
| 3523 | + arr = fp.read(read_length) |
| 3524 | + current_pos += read_length |
| 3525 | + self.assertEqual(fp.tell(), current_pos) |
| 3526 | + self.assertEqual(arr, txt[current_pos - read_length:current_pos]) |
| 3527 | + self.assertEqual(fp._left, fp._compress_left) |
| 3528 | + read_count = sio.bytes_read - old_count |
| 3529 | + self.assertLessEqual(read_count, read_buffer_size) |
| 3530 | + |
| 3531 | + # eof flags test |
| 3532 | + fp.seek(0, os.SEEK_END) |
| 3533 | + fp.seek(12345, os.SEEK_SET) |
| 3534 | + current_pos = 12345 |
| 3535 | + arr = fp.read(read_length) |
| 3536 | + current_pos += read_length |
| 3537 | + self.assertEqual(arr, txt[current_pos - read_length:current_pos]) |
| 3538 | + |
| 3539 | + |
3457 | 3540 | if __name__ == "__main__":
|
3458 | 3541 | unittest.main()
|
0 commit comments