source: cows_wps/trunk/cows_wps/process_handler/fileset.py @ 7018

"""
Implements the data structure originally designed by Ag to store a nested set of
data or zipfiles.

"""
import logging
import zipfile
import os
import stat

class FLAG:
    ROOT = 'r'
    DATA = 'data'
    ZIP = 'zip'
    IMAGE = 'image'


log = logging.getLogger(__name__)

class FileSetEntry(object):
    """
    @ivar flag: A flag designating the file type
    @ivar size: The size of the file
    @ivar name: the name of the file relative to it's container
    @ivar info: Some description text about the file
    @ivar capabilities: Some text used to define what can be done with the file
    
    """
    def __init__(self, flag, name, size, info):
        self.flag = flag
        self.name = name
        self.size = size
        self.info = info
            
    def __get_capabilities(self):
        file_capabilities = ""
        
        if self.flag == FLAG.ZIP:
            file_capabilities += "download expand_description"
        
        return file_capabilities
    
    capabilities = property(__get_capabilities, None, None, "")    

class FileSet(FileSetEntry):
    """
    An aggregation of FileSetEntry objects.  FileSet inherits from FileSetEntry 
    so that FileSet objects can contain other FileSet objects (typically 
    representing zipfiles).
    
    @ivar contents: A mutable sequence of FileSet contents
        
    """
        
    def __init__(self, flag=FLAG.ROOT, name='root', size=None, info="Root fileset"):
        super(FileSet, self).__init__(flag, name, size, info)
        self.contents = []
        
    def __iter__(self):
        return iter(self.contents)
        
def dryRunZipFileSet(fileSet, context, maxFileVolume):
    """
    Pretend to zip up the contents of a FileSet into zipfiles, guessing file volumes.
    
    """
    # Compression ratios are taken from context.config.
    compressionRatios = dict(context.config['compression_ratios'])
    #log.debug("compressionRatios: %s" % compressionRatios)
    
    i = 1
    size = 0
    zipFileSets = []
    
    for fs in fileSet.contents:
        
        
        # New zipFile
        if size == 0:
            zipname = 'output_%d.zip' % i
            zipFileSet = FileSet(FLAG.ZIP, name=zipname, info='Zip file %d' % i)

        ext = os.path.splitext(fs.name)[1]
        ext = ext.strip('.')

        dratio = float(compressionRatios.get('default_ratio', 1.0))
        ratio = float(compressionRatios.get(ext, dratio))
        
        size += fs.size * ratio
        zipFileSet.contents.append(fs)
        
        if size >= maxFileVolume:
            zipFileSet.size = int(size)
            
            log.debug("total file size for zipFileSet.name = %s is %s MB" % (zipFileSet.name, float(zipFileSet.size)/4**10))
            
            zipFileSets.append(zipFileSet)
            size = 0
            i += 1
    
    #update the last set with the size, only if there have been more files added        
    if size > 0:
        zipFileSet.size = int(size)
        zipFileSets.append(zipFileSet)
        log.debug("total file size for zipFileSet.name = %s is %s MB" % (zipFileSet.name, float(zipFileSet.size)/4**10))
    
    if len(zipFileSets) == 1:
        zipFileSets[0].name = 'output.zip'
    
    fileSet.contents = zipFileSets
    
    return fileSet

def zipFileSet(fileSet, context, maxFileVolume, keep = False):
    """
    Zip up the contents of a FileSet into zipfiles of maximum size 
    I{maxFileVolume}.
    
    @return: I{fileSet} mutated to contain the zipfiles created.
    
    """
    i = 1
    size = 0
    zipFileSets = []
    
    
    for fs in fileSet.contents:
        # New zipfile
        if size == 0:
            zipname = 'output_%d.zip' % i
            fullZipFilename = os.path.join(context.outputDir, zipname)
            zip = zipfile.ZipFile(fullZipFilename, 'w', zipfile.ZIP_DEFLATED)
            zipFileSet = FileSet(flag=FLAG.ZIP, name=zipname, info='Zip file %d' % i)
        
        #!TODO: This way of invoking zip will not correctly archive directories.
        zip.write(os.path.join(context.outputDir, fs.name), fs.name)
        zi = zip.getinfo(fs.name)
        size += zi.compress_size
        
        zi.comment = fs.info
        zipFileSet.contents.append(fs)
        
        if size >= maxFileVolume:
            # the true file size is slightly larger than the compressed size
            # of all the contained files, query the os to get the true size of
            # the zip file. Also need to close the zip file first.
            filename = zip.filename
            zip.close()
            zipFileSet.size = _getFileSize(filename)
            zipFileSets.append(zipFileSet)
            size = 0
            i += 1
    
    filename = zip.filename
    zip.close()
    zipFileSet.size = _getFileSize(filename)
    zipFileSets.append(zipFileSet)
    
    log.debug("len(zipFileSets) = %s" % (len(zipFileSets),))
    
    if len(zipFileSets) == 1:
        log.debug("renaming")
        oldName = 'output_1.zip'
        newName = 'output.zip'
        
        os.rename(os.path.join(context.outputDir, oldName), 
                  os.path.join(context.outputDir, newName))
        
        zipFileSets[0].name = newName
        
    
    log.info("Output zipped up into %s zip files" % (len(zipFileSets),))
    
    # Now we know this zipfile is complete clean up.
    if keep == False:
        for zipFileSet in zipFileSets:
            for fs in zipFileSet.contents:
                os.remove(os.path.join(context.outputDir, fs.name))
    
    fileSet.contents = zipFileSets
    return fileSet


def _getFileSize(fullFileName):
    fileStats = os.stat(fullFileName)
    return fileStats[stat.ST_SIZE]