| import os |
| import io |
| import contextlib |
| import tempfile |
| import shutil |
| import errno |
| import zipfile |
| |
| |
| @contextlib.contextmanager |
| def tempdir(): |
| """Create a temporary directory in a context manager.""" |
| td = tempfile.mkdtemp() |
| try: |
| yield td |
| finally: |
| shutil.rmtree(td) |
| |
| |
| def mkdir_p(*args, **kwargs): |
| """Like `mkdir`, but does not raise an exception if the |
| directory already exists. |
| """ |
| try: |
| return os.mkdir(*args, **kwargs) |
| except OSError as exc: |
| if exc.errno != errno.EEXIST: |
| raise |
| |
| |
| def dir_to_zipfile(root): |
| """Construct an in-memory zip file for a directory.""" |
| buffer = io.BytesIO() |
| zip_file = zipfile.ZipFile(buffer, 'w') |
| for root, dirs, files in os.walk(root): |
| for path in dirs: |
| fs_path = os.path.join(root, path) |
| rel_path = os.path.relpath(fs_path, root) |
| zip_file.writestr(rel_path + '/', '') |
| for path in files: |
| fs_path = os.path.join(root, path) |
| rel_path = os.path.relpath(fs_path, root) |
| zip_file.write(fs_path, rel_path) |
| return zip_file |
