Write data from stdin to a ZIP file on the fly

datafile.py

#!/usr/bin/env python3

from collections import deque as Deque
from zipstream   import ZipFile

class Buffer:
	def __init__(self):
		self.queue = Deque()
		self.finished = False

	def __bool__(self):
		return bool(self.queue)

	def __iter__(self):
		return self

	def __next__(self):
		if self.finished:
			raise StopIteration
		if self.queue:
			return self.queue.popleft()
		else:
			return b''

	def append(self, chunk):
		assert not self.finished
		self.queue.append(chunk)

	def finish(self):
		self.finished = True


class DataFile(object):
	def __init__(self, fileobj):
		self.fp = fileobj
		self.buffer = Buffer()

		zip = ZipFile() # sic
		zip.write_iter('header.txt', [b'This is a header\n'])
		zip.write_iter('data.bin', self.buffer)
		zip.write_iter('trailer.txt', [b'This is a trailer\n'])
		self.zipiter = iter(zip)

	def __enter__(self):
		return self

	def __exit__(self, exc_type, exc_value, exc_traceback):
		self.close()

	def write(self, chunk):
		self.buffer.append(chunk)
		while self.buffer:
			self.fp.write(next(self.zipiter))

	def close(self):
		try:
			self.buffer.finish()
			while True:
				try:
					self.fp.write(next(self.zipiter))
				except StopIteration:
					break
		finally:
			self.fp.close()
			self.fp = None


if __name__ == '__main__':
	from sys import argv, stdin

	with DataFile(open(argv[1], 'wb')) as df:
		while True:
			chunk = stdin.buffer.read1()
			if not chunk: break
			df.write(chunk)