pakfire/compress.py

   1 #!/usr/bin/python
   2 ###############################################################################
   3 #                                                                             #
   4 # Pakfire - The IPFire package management system                              #
   5 # Copyright (C) 2011 Pakfire development team                                 #
   6 #                                                                             #
   7 # This program is free software: you can redistribute it and/or modify        #
   8 # it under the terms of the GNU General Public License as published by        #
   9 # the Free Software Foundation, either version 3 of the License, or           #
  10 # (at your option) any later version.                                         #
  11 #                                                                             #
  12 # This program is distributed in the hope that it will be useful,             #
  13 # but WITHOUT ANY WARRANTY; without even the implied warranty of              #
  14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
  15 # GNU General Public License for more details.                                #
  16 #                                                                             #
  17 # You should have received a copy of the GNU General Public License           #
  18 # along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
  19 #                                                                             #
  20 ###############################################################################
  21
  22 import lzma
  23 import os
  24 import progressbar
  25 import zlib
  26
  27 from constants import *
  28 from i18n import _
  29
  30 PROGRESS_WIDGETS = [
  31         progressbar.Bar(left="[", right="]"),
  32         "  ",
  33         progressbar.Percentage(),
  34         "  ",
  35         progressbar.ETA(),
  36         "  ",
  37 ]
  38
  39 def __compress_helper(i, o, comp, flush, progress=None):
  40         if progress:
  41                 widgets = [ "%-30s  " % os.path.basename(filename)] + PROGRESS_WIDGETS
  42
  43                 maxval = os.path.getsize(filename)
  44
  45                 progress = progressbar.ProgressBar(
  46                         widgets=widgets,
  47                         maxval=maxval,
  48                 )
  49
  50                 progress.start()
  51
  52         size = 0
  53         buf = i.read(BUFFER_SIZE)
  54         while buf:
  55                 if progress:
  56                         size += len(buf)
  57                         progress.update(size)
  58
  59                 o.write(comp(buf))
  60
  61                 buf = i.read(BUFFER_SIZE)
  62
  63         o.write(flush())
  64
  65         if progress:
  66                 progress.finish()
  67
  68 def compress(filename, filename2=None, algo="xz", progress=None):
  69         i = open(filename)
  70
  71         if not filename2:
  72                 filename2 = filename
  73                 os.unlink(filename)
  74
  75         o = open(filename2, "w")
  76
  77         compressobj(i, o, algo="xz", progress=None)
  78
  79         i.close()
  80         o.close()
  81
  82 def compressobj(i, o, algo="xz", progress=None):
  83         comp = None
  84         if algo == "xz":
  85                 comp = lzma.LZMACompressor()
  86
  87         elif algo == "zlib":
  88                 comp = zlib.compressobj(9)
  89
  90         return __compress_helper(i, o, comp.compress, comp.flush, progress=progress)
  91
  92 def decompress(filename, filename2=None, algo="xz", progress=None):
  93         i = open(filename)
  94
  95         if not filename2:
  96                 filename2 = filename
  97                 os.unlink(filename)
  98
  99         o = open(filename2, "w")
 100
 101         decompressobj(i, o, algo="xz", progress=None)
 102
 103         i.close()
 104         o.close()
 105
 106 def decompressobj(i, o, algo="xz", progress=None):
 107         comp = None
 108         if algo == "xz":
 109                 comp = lzma.LZMADecompressor()
 110
 111         elif algo == "zlib":
 112                 comp = zlib.decompressobj(9)
 113
 114         return __compress_helper(i, o, comp.decompress, comp.flush, progress=progress)
 115
 116 def compress_file(inputfile, outputfile, message="", algo="xz", progress=True):
 117         """
 118                 Compress a file in place.
 119         """
 120         assert os.path.exists(inputfile)
 121
 122         # Get total size of the file for the progressbar.
 123         total_size = os.path.getsize(inputfile)
 124
 125         # Open the input file for reading.
 126         i = open(inputfile, "r")
 127
 128         # Open the output file for wrinting.
 129         o = open(outputfile, "w")
 130
 131         if progress:
 132                 if not message:
 133                         message = _("Compressing %s") % os.path.basename(filename)
 134
 135                 progress = progressbar.ProgressBar(
 136                         widgets = ["%-40s" % message, " ",] + PROGRESS_WIDGETS,
 137                         maxval = total_size,
 138                 )
 139
 140                 progress.start()
 141
 142         if algo == "xz":
 143                 compressor = lzma.LZMACompressor()
 144         elif algo == "zlib":
 145                 comp = zlib.decompressobj(9)
 146         else:
 147                 raise Exception, "Unknown compression choosen: %s" % algo
 148
 149         size = 0
 150         while True:
 151                 buf = i.read(BUFFER_SIZE)
 152                 if not buf:
 153                         break
 154
 155                 # Update progressbar.
 156                 size += len(buf)
 157                 if progress:
 158                         progress.update(size)
 159
 160                 # Compress the bits in buf.
 161                 buf = compressor.compress(buf)
 162
 163                 # Write the compressed output.
 164                 o.write(buf)
 165
 166         # Flush all buffers.
 167         buf = compressor.flush()
 168         o.write(buf)
 169
 170         # Close the progress bar.
 171         if progress:
 172                 progress.finish()
 173
 174         i.close()
 175         o.close()