|
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