|
4 | 4 | from collections import abc |
5 | 5 | import dataclasses |
6 | 6 | import gzip |
7 | | -from io import BufferedIOBase, BytesIO, RawIOBase, TextIOWrapper |
| 7 | +from io import BufferedIOBase, BytesIO, RawIOBase, StringIO, TextIOWrapper |
8 | 8 | import mmap |
9 | 9 | import os |
10 | | -from typing import IO, Any, AnyStr, Dict, List, Mapping, Optional, Tuple, cast |
| 10 | +from typing import IO, Any, AnyStr, Dict, List, Mapping, Optional, Tuple, Union, cast |
11 | 11 | from urllib.parse import ( |
12 | 12 | urljoin, |
13 | 13 | urlparse as parse_url, |
@@ -713,17 +713,36 @@ def __init__( |
713 | 713 | archive_name: Optional[str] = None, |
714 | 714 | **kwargs, |
715 | 715 | ): |
716 | | - if mode in ["wb", "rb"]: |
717 | | - mode = mode.replace("b", "") |
| 716 | + mode = mode.replace("b", "") |
718 | 717 | self.archive_name = archive_name |
| 718 | + self.multiple_write_buffer: Optional[Union[StringIO, BytesIO]] = None |
| 719 | + |
719 | 720 | kwargs_zip: Dict[str, Any] = {"compression": zipfile.ZIP_DEFLATED} |
720 | 721 | kwargs_zip.update(kwargs) |
| 722 | + |
721 | 723 | super().__init__(file, mode, **kwargs_zip) # type: ignore[arg-type] |
722 | 724 |
|
723 | 725 | def write(self, data): |
| 726 | + # buffer multiple write calls, write on flush |
| 727 | + if self.multiple_write_buffer is None: |
| 728 | + self.multiple_write_buffer = ( |
| 729 | + BytesIO() if isinstance(data, bytes) else StringIO() |
| 730 | + ) |
| 731 | + self.multiple_write_buffer.write(data) |
| 732 | + |
| 733 | + def flush(self) -> None: |
| 734 | + # write to actual handle and close write buffer |
| 735 | + if self.multiple_write_buffer is None or self.multiple_write_buffer.closed: |
| 736 | + return |
| 737 | + |
724 | 738 | # ZipFile needs a non-empty string |
725 | 739 | archive_name = self.archive_name or self.filename or "zip" |
726 | | - super().writestr(archive_name, data) |
| 740 | + with self.multiple_write_buffer: |
| 741 | + super().writestr(archive_name, self.multiple_write_buffer.getvalue()) |
| 742 | + |
| 743 | + def close(self): |
| 744 | + self.flush() |
| 745 | + super().close() |
727 | 746 |
|
728 | 747 | @property |
729 | 748 | def closed(self): |
|
0 commit comments