45 lines
1.1 KiB
Python
45 lines
1.1 KiB
Python
|
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
|