Source code for galaxy.util.streamball

"""
A simple wrapper for writing tarballs as a stream.
"""
from __future__ import absolute_import

import logging
import os
import tarfile

from galaxy.exceptions import ObjectNotFound
from .path import safe_walk

log = logging.getLogger(__name__)


[docs]class StreamBall(object):
[docs] def __init__(self, mode, members=None): self.members = members if members is None: self.members = [] self.mode = mode self.wsgi_status = None self.wsgi_headeritems = None
[docs] def add(self, file, relpath, check_file=False): if check_file and len(file) > 0: if not os.path.isfile(file): raise ObjectNotFound else: self.members.append((file, relpath)) else: self.members.append((file, relpath))
[docs] def stream(self, environ, start_response): response_write = start_response(self.wsgi_status, self.wsgi_headeritems) class tarfileobj(object): def write(self, *args, **kwargs): response_write(*args, **kwargs) tf = tarfile.open(mode=self.mode, fileobj=tarfileobj()) for (file, rel) in self.members: tf.add(file, arcname=rel) tf.close() return []
[docs]class ZipBall(object):
[docs] def __init__(self, tmpf, tmpd): self._tmpf = tmpf self._tmpd = tmpd
[docs] def stream(self, environ, start_response): response_write = start_response(self.wsgi_status, self.wsgi_headeritems) tmpfh = open(self._tmpf) response_write(tmpfh.read()) tmpfh.close() try: os.unlink(self._tmpf) os.rmdir(self._tmpd) except OSError: log.exception("Unable to remove temporary library download archive and directory") return []
[docs]def stream_archive(trans, path, upstream_gzip=False): archive_type_string = 'w|gz' archive_ext = 'tgz' if upstream_gzip: archive_type_string = 'w|' archive_ext = 'tar' archive = StreamBall(mode=archive_type_string) for root, directories, files in safe_walk(path): for filename in files: p = os.path.join(root, filename) relpath = os.path.relpath(p, os.path.join(path, os.pardir)) archive.add(file=os.path.join(path, p), relpath=relpath) archive_name = "%s.%s" % (os.path.basename(path), archive_ext) trans.response.set_content_type("application/x-tar") trans.response.headers["Content-Disposition"] = 'attachment; filename="{}"'.format(archive_name) archive.wsgi_status = trans.response.wsgi_status() archive.wsgi_headeritems = trans.response.wsgi_headeritems() return archive.stream