summaryrefslogtreecommitdiffstats
path: root/utils/log_picker/archiving.py
blob: fb3da9c392112b066d0daf5674f5c0b627a644e5 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import os
import tempfile
import tarfile
import bz2


class ArchivationError(Exception):
    pass

class NoFilesArchivationError(ArchivationError):
    pass

class ArchiveBaseClass(object):
    """Base class for archive classes."""

    _compression = False
    _ext = ".ext"
    _mimetype = ""
    
    def __init__(self, *args, **kwargs):
        self._tar_ext = ".tar"
    
    @property
    def support_compression(self):
        """Return True if compression is supported/used."""
        return self._compression
        
    @property
    def file_ext(self):
        """Return extension for output file."""
        return self._ext
       
    @property 
    def mimetype(self):
        """Return archive mime type."""
        return self._mimetype
    
    def _create_tmp_tar(self, filelist):
        _, tmpfile = tempfile.mkstemp(suffix=self._tar_ext)
        tar = tarfile.open(tmpfile, "w")
        for name in filelist:
            pieces = name.rsplit('/', 2)
            arcname = "%s/%s" % (pieces[-2], pieces[-1])
            tar.add(name, arcname=arcname)
        tar.close()
        return tmpfile
       
    def create_archive(self, outfilename, filelist):
        raise NotImplementedError()


class Bzip2Archive(ArchiveBaseClass):
    """Class for bzip2 compression."""
    
    _compression = True
    _ext = ".bz2"
    _mimetype = "application/x-bzip2"
    
    def __init__(self, usetar=True, *args, **kwargs):
        ArchiveBaseClass.__init__(self, args, kwargs)
        self.usetar = usetar

    @property
    def file_ext(self):
        """Return extension for output file."""
        if self.usetar:
            return "%s%s" % (self._tar_ext, self._ext)
        return self._ext

    def create_archive(self, outfilename, filelist):
        """Create compressed archive containing files listed in filelist."""
        if not filelist:
            raise NoFilesArchivationError("No files to archive.")
        
        size = 0
        for file in filelist:
            size += os.path.getsize(file)
        if size <= 0:
            raise NoFilesArchivationError("No files to archive.")
            
        if not self.usetar and len(filelist) > 1:
            raise ArchivationError( \
                            "Bzip2 cannot archive multiple files without tar.")
        
        if self.usetar:
            f_in_path = self._create_tmp_tar(filelist)
        else:
            f_in_path = filelist[0]
         
        f_in = open(f_in_path, 'rb')
        f_out = bz2.BZ2File(outfilename, 'w')
        f_out.writelines(f_in)
        f_out.close()
        f_in.close()
        
        if self.usetar:
            os.remove(f_in_path)
        
        return outfilename