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