zipHelpers.py

from io import BytesIO
import zipfile

class InMemoryZip(object):
	def __init__(self, items = {}):
		# Create the in-memory file-like object for working w/IMZ
		self.in_memory_zip = BytesIO()
		self.items: "dict[str, bytes]" = {}
		for i in items:
			self.append(i, items[i])
	# Just zip it, zip it
	def append(self, filename_in_zip: str, file_contents: bytes):
		# Record this file.
		self.items[filename_in_zip] = file_contents
		# Appends a file with name filename_in_zip and contents of
		# file_contents to the in-memory zip.
		# Get a handle to the in-memory zip in append mode
		zf = zipfile.ZipFile(self.in_memory_zip, "a", zipfile.ZIP_DEFLATED, False)
		# Write the file to the in-memory zip
		zf.writestr(filename_in_zip, file_contents)
		# Mark the files as having been created on Windows so that
		# Unix permissions are not inferred as 0000
		for zfile in zf.filelist:
			zfile.create_system = 0
		return self
	def append_multiple(self, files: "dict[str, bytes]"):
		# Appends a file with name filename_in_zip and contents of
		# file_contents to the in-memory zip.
		# Get a handle to the in-memory zip in append mode
		zf = zipfile.ZipFile(self.in_memory_zip, "a", zipfile.ZIP_DEFLATED, False)
		# Record this file.
		for filename_in_zip in files:
			file_contents = files[filename_in_zip]
			self.items[filename_in_zip] = file_contents
			# Write the file to the in-memory zip
			zf.writestr(filename_in_zip, file_contents)
		# Mark the files as having been created on Windows so that
		# Unix permissions are not inferred as 0000
		for zfile in zf.filelist:
			zfile.create_system = 0
		return self
	def read(self):
		# Returns a string with the contents of the in-memory zip.
		self.in_memory_zip.seek(0)
		return self.in_memory_zip.read()
	# Zip it, zip it, zip it
	def writetofile(self, filename):
		# Writes the in-memory zip to a physical file.
		with open(filename, "wb") as file:
			file.write(self.read())
def extract_zip(input_zip):
	input_zip=zipfile.ZipFile(input_zip)
	imz = InMemoryZip({name: input_zip.read(name) for name in input_zip.namelist()})
	return imz
if __name__ == "__main__":
	# Run a test
	imz = InMemoryZip()
	imz.append("testfile.txt", "Make a test").append("hi/testfile2.txt", "And another one")
	imz.writetofile("testfile.zip")
	print("testfile.zip created")
	imz = extract_zip("testfile.zip")
	print(imz.items)