1import os 2import io 3import contextlib 4import tempfile 5import shutil 6import errno 7import zipfile 8 9 10@contextlib.contextmanager 11def tempdir(): 12 """Create a temporary directory in a context manager.""" 13 td = tempfile.mkdtemp() 14 try: 15 yield td 16 finally: 17 shutil.rmtree(td) 18 19 20def mkdir_p(*args, **kwargs): 21 """Like `mkdir`, but does not raise an exception if the 22 directory already exists. 23 """ 24 try: 25 return os.mkdir(*args, **kwargs) 26 except OSError as exc: 27 if exc.errno != errno.EEXIST: 28 raise 29 30 31def dir_to_zipfile(root): 32 """Construct an in-memory zip file for a directory.""" 33 buffer = io.BytesIO() 34 zip_file = zipfile.ZipFile(buffer, 'w') 35 for root, dirs, files in os.walk(root): 36 for path in dirs: 37 fs_path = os.path.join(root, path) 38 rel_path = os.path.relpath(fs_path, root) 39 zip_file.writestr(rel_path + '/', '') 40 for path in files: 41 fs_path = os.path.join(root, path) 42 rel_path = os.path.relpath(fs_path, root) 43 zip_file.write(fs_path, rel_path) 44 return zip_file 45