| 20 | |
| 21 | |
| 22 | class InMemoryZipFile: |
| 23 | def __init__( |
| 24 | self, file_name: Optional[str | Path] = None, compression: int = zipfile.ZIP_DEFLATED, debug: int = 0 |
| 25 | ) -> None: |
| 26 | # Create the in-memory file-like object |
| 27 | self._file_name: Optional[str | Path] = str(file_name) if hasattr(file_name, "_from_parts") else file_name |
| 28 | self.in_memory_data = BytesIO() |
| 29 | # Create the in-memory zipfile |
| 30 | self.in_memory_zip = zipfile.ZipFile(self.in_memory_data, "w", compression, True) |
| 31 | self.in_memory_zip.debug = debug |
| 32 | |
| 33 | def writestr(self, filename_in_zip: str | zipfile.ZipInfo, file_contents: bytes | str) -> None: |
| 34 | """Appends a file with name filename_in_zip and contents of |
| 35 | file_contents to the in-memory zip.""" |
| 36 | self.in_memory_zip.writestr(filename_in_zip, file_contents) |
| 37 | |
| 38 | def write_to_file(self, filename: str | bytes | PathLike[str] | PathLike[bytes] | int) -> None: |
| 39 | """Writes the in-memory zip to a file.""" |
| 40 | # Mark the files as having been created on Windows so that |
| 41 | # Unix permissions are not inferred as 0000 |
| 42 | for zfile in self.in_memory_zip.filelist: |
| 43 | zfile.create_system = 0 |
| 44 | self.in_memory_zip.close() |
| 45 | with open(filename, "wb") as f: |
| 46 | f.write(self.data) |
| 47 | |
| 48 | @property |
| 49 | def data(self) -> bytes: |
| 50 | return self.in_memory_data.getvalue() |
| 51 | |
| 52 | def __enter__(self) -> InMemoryZipFile: |
| 53 | return self |
| 54 | |
| 55 | def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: |
| 56 | if self._file_name: |
| 57 | self.write_to_file(self._file_name) |