/lib/galaxy/util/streamball.py

 1"""
 2A simple wrapper for writing tarballs as a stream.
 3"""
 4import os
 5import logging, tarfile
 6from galaxy.exceptions import ObjectNotFound
 7
 8log = logging.getLogger( __name__ )
 9
10class StreamBall( object ):
11    def __init__( self, mode, members=None ):
12        self.members = members
13        if members is None:
14            self.members = {}
15        self.mode = mode
16        self.wsgi_status = None
17        self.wsgi_headeritems = None
18    def add( self, file, relpath, check_file=False):
19        if check_file and len(file)>0:
20            if not os.path.isfile(file):
21                raise ObjectNotFound
22        else:
23            self.members[file] = relpath
24    def stream( self, environ, start_response ):
25        response_write = start_response( self.wsgi_status, self.wsgi_headeritems )
26        class tarfileobj:
27            def write( self, *args, **kwargs ):
28                response_write( *args, **kwargs )
29        tf = tarfile.open( mode=self.mode, fileobj=tarfileobj() )
30        for file, rel in self.members.items():
31            tf.add( file, arcname=rel )
32        tf.close()
33        return []
34
35class ZipBall(object):
36    def __init__(self, tmpf, tmpd):
37        self._tmpf = tmpf
38        self._tmpd = tmpd
39    def stream(self, environ, start_response):
40        response_write = start_response( self.wsgi_status, self.wsgi_headeritems )
41        tmpfh = open( self._tmpf )
42        response_write(tmpfh.read())
43        tmpfh.close()
44        try:
45            os.unlink( self._tmpf )
46            os.rmdir( self._tmpd )
47        except OSError:
48            log.exception( "Unable to remove temporary library download archive and directory" )
49        return []