|
9 | 9 | import zlib
|
10 | 10 | import builtins
|
11 | 11 | import io
|
| 12 | +import _pyio as pyio |
12 | 13 | import _compression
|
| 14 | +import errno |
13 | 15 |
|
14 | 16 | __all__ = ["BadGzipFile", "GzipFile", "open", "compress", "decompress"]
|
15 | 17 |
|
@@ -445,6 +447,115 @@ def _read_gzip_header(fp):
|
445 | 447 | _read_exact(fp, 2) # Read & discard the 16-bit header CRC
|
446 | 448 | return last_mtime
|
447 | 449 |
|
| 450 | +class _Writable(): |
| 451 | + def writable(self): |
| 452 | + return True |
| 453 | + |
| 454 | +class GzipWriter(pyio.BufferedWriter): |
| 455 | + def __init__(self, fp, compresslevel=_COMPRESS_LEVEL_BEST, buffer_size=pyio.DEFAULT_BUFFER_SIZE, mtime=None): |
| 456 | + super().__init__(fp, buffer_size=buffer_size) |
| 457 | + |
| 458 | + self.fileobj = fp |
| 459 | + self.crc = zlib.crc32(b"") |
| 460 | + self.size = 0 |
| 461 | + self.writebuf = [] |
| 462 | + self.bufsize = 0 |
| 463 | + self.offset = 0 # Current file offset for seek(), tell(), etc |
| 464 | + self.compress = zlib.compressobj(compresslevel, |
| 465 | + zlib.DEFLATED, |
| 466 | + -zlib.MAX_WBITS, |
| 467 | + zlib.DEF_MEM_LEVEL, |
| 468 | + 0) |
| 469 | + self._write_mtime = mtime |
| 470 | + |
| 471 | + self._write_gzip_header(compresslevel) |
| 472 | + |
| 473 | + def close(self): |
| 474 | + fileobj = self.fileobj |
| 475 | + self.flush() |
| 476 | + fileobj.write(self.compress.flush()) |
| 477 | + write32u(fileobj, self.crc) |
| 478 | + # self.size may exceed 2 GiB, or even 4 GiB |
| 479 | + write32u(fileobj, self.size & 0xffffffff) |
| 480 | + self.fileobj.close() |
| 481 | + |
| 482 | + def write(self, data): |
| 483 | + super().write(data) |
| 484 | + |
| 485 | + def _flush_unlocked(self): |
| 486 | + if self.closed: |
| 487 | + raise ValueError("flush on closed file") |
| 488 | + while self._write_buf: |
| 489 | + try: |
| 490 | + #n = self.raw.write(self._write_buf) |
| 491 | + n = self.compress_and_write(self._write_buf) |
| 492 | + except BlockingIOError: |
| 493 | + raise RuntimeError("self.raw should implement RawIOBase: it " |
| 494 | + "should not raise BlockingIOError") |
| 495 | + if n is None: |
| 496 | + raise BlockingIOError( |
| 497 | + errno.EAGAIN, |
| 498 | + "write could not complete without blocking", 0) |
| 499 | + if n > len(self._write_buf) or n < 0: |
| 500 | + raise OSError("write() returned incorrect number of bytes") |
| 501 | + del self._write_buf[:n] |
| 502 | + |
| 503 | + def compress_and_write(self,data): |
| 504 | + self._check_not_closed() |
| 505 | + if self.fileobj is None: |
| 506 | + raise ValueError("write() on closed GzipFile object") |
| 507 | + |
| 508 | + if isinstance(data, (bytes, bytearray)): |
| 509 | + length = len(data) |
| 510 | + else: |
| 511 | + # accept any data that supports the buffer protocol |
| 512 | + data = memoryview(data) |
| 513 | + length = data.nbytes |
| 514 | + |
| 515 | + if length > 0: |
| 516 | + self.fileobj.write(self.compress.compress(data)) |
| 517 | + self.size += length |
| 518 | + self.crc = zlib.crc32(data, self.crc) |
| 519 | + self.offset += length |
| 520 | + |
| 521 | + return length |
| 522 | + |
| 523 | + def _write_gzip_header(self, compresslevel): |
| 524 | + self.fileobj.write(b'\037\213') # magic header |
| 525 | + self.fileobj.write(b'\010') # compression method |
| 526 | + try: |
| 527 | + # RFC 1952 requires the FNAME field to be Latin-1. Do not |
| 528 | + # include filenames that cannot be represented that way. |
| 529 | + fname = os.path.basename(self.name) |
| 530 | + if not isinstance(fname, bytes): |
| 531 | + fname = fname.encode('latin-1') |
| 532 | + if fname.endswith(b'.gz'): |
| 533 | + fname = fname[:-3] |
| 534 | + except UnicodeEncodeError: |
| 535 | + fname = b'' |
| 536 | + flags = 0 |
| 537 | + if fname: |
| 538 | + flags = FNAME |
| 539 | + self.fileobj.write(chr(flags).encode('latin-1')) |
| 540 | + mtime = self._write_mtime |
| 541 | + if mtime is None: |
| 542 | + mtime = time.time() |
| 543 | + write32u(self.fileobj, int(mtime)) |
| 544 | + if compresslevel == _COMPRESS_LEVEL_BEST: |
| 545 | + xfl = b'\002' |
| 546 | + elif compresslevel == _COMPRESS_LEVEL_FAST: |
| 547 | + xfl = b'\004' |
| 548 | + else: |
| 549 | + xfl = b'\000' |
| 550 | + self.fileobj.write(xfl) |
| 551 | + self.fileobj.write(b'\377') |
| 552 | + if fname: |
| 553 | + self.fileobj.write(fname + b'\000') |
| 554 | + |
| 555 | + def _check_not_closed(self): |
| 556 | + return self.fileobj.closed |
| 557 | + |
| 558 | + |
448 | 559 |
|
449 | 560 | class _GzipReader(_compression.DecompressReader):
|
450 | 561 | def __init__(self, fp):
|
|
0 commit comments