mic/archive.py

   1 #
   2 # Copyright (c) 2013 Intel Inc.
   3 #
   4 # This program is free software; you can redistribute it and/or modify it
   5 # under the terms of the GNU General Public License as published by the Free
   6 # Software Foundation; version 2 of the License
   7 #
   8 # This program is distributed in the hope that it will be useful, but
   9 # WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  10 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  11 # for more details.
  12 #
  13 # You should have received a copy of the GNU General Public License along
  14 # with this program; if not, write to the Free Software Foundation, Inc., 59
  15 # Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  16
  17 # Following messages should be disabled in pylint
  18 #  * Used * or ** magic (W0142)
  19 #  * Unused variable (W0612)
  20 #  * Used builtin function (W0141)
  21 #  * Invalid name for type (C0103)
  22 #  * Popen has no '%s' member (E1101)
  23 # pylint: disable=W0142, W0612, W0141, C0103, E1101
  24
  25 """ Compression and Archiving
  26
  27 Utility functions for creating archive files (tarballs, zip files, etc)
  28 and compressing files (gzip, bzip2, lzop, etc)
  29 """
  30
  31 import os
  32 import shutil
  33 import tempfile
  34 import subprocess
  35 from mic import msger
  36
  37 __all__ = [
  38             "get_compress_formats",
  39             "compress",
  40             "decompress",
  41             "get_archive_formats",
  42             "get_archive_suffixes",
  43             "make_archive",
  44             "extract_archive",
  45             "compressing",
  46             "packing",
  47           ]
  48
  49 # TODO: refine Error class for archive/extract
  50
  51 def which(binary, path=None):
  52     """ Find 'binary' in the directories listed in 'path'
  53
  54     @binary: the executable file to find
  55     @path: the suposed path to search for, use $PATH if None
  56     @retval: the absolute path if found, otherwise None
  57     """
  58     if path is None:
  59         path = os.environ["PATH"]
  60     for apath in path.split(os.pathsep):
  61         fpath = os.path.join(apath, binary)
  62         if os.path.isfile(fpath) and os.access(fpath, os.X_OK):
  63             return fpath
  64     return None
  65
  66 def _call_external(cmdln_or_args):
  67     """ Wapper for subprocess calls.
  68
  69     @cmdln_or_args: command line to be joined before execution.
  70     @retval: a tuple (returncode, outdata, errdata).
  71     """
  72     if isinstance(cmdln_or_args, list):
  73         shell = False
  74     else:
  75         shell = True
  76     msger.info("Running command: " + " ".join(cmdln_or_args))
  77
  78     proc = subprocess.Popen(cmdln_or_args, shell=shell,
  79                             stdout=subprocess.PIPE,
  80                             stderr=subprocess.STDOUT)
  81     (outdata, errdata) = proc.communicate()
  82
  83     return (proc.returncode, outdata, errdata)
  84
  85 def _do_gzip(input_name, compression=True):
  86     """ Compress/decompress the file with 'gzip' utility.
  87
  88     @input_name: the file name to compress/decompress
  89     @compress: True for compressing, False for decompressing
  90     @retval: the path of the compressed/decompressed file
  91     """
  92     if which("pigz") is not None:
  93         compressor = "pigz"
  94     else:
  95         compressor = "gzip"
  96
  97     if compression:
  98         cmdln = [compressor, "-f", input_name]
  99     else:
 100         cmdln = [compressor, "-d", "-f", input_name]
 101
 102     _call_external(cmdln)
 103
 104     if compression:
 105         output_name = input_name + ".gz"
 106     else:
 107         # suppose that file name is suffixed with ".gz"
 108         output_name = os.path.splitext(input_name)[0]
 109
 110     return output_name
 111
 112 def _do_bzip2(input_name, compression=True):
 113     """ Compress/decompress the file with 'bzip2' utility.
 114
 115     @input_name: the file name to compress/decompress
 116     @compress: True for compressing, False for decompressing
 117     @retval: the path of the compressed/decompressed file
 118     """
 119     if which("pbzip2") is not None:
 120         compressor = "pbzip2"
 121     else:
 122         compressor = "bzip2"
 123
 124     if compression:
 125         cmdln = [compressor, "-f", input_name]
 126     else:
 127         cmdln = [compressor, "-d", "-f", input_name]
 128
 129     _call_external(cmdln)
 130
 131     if compression:
 132         output_name = input_name + ".bz2"
 133     else:
 134         # suppose that file name is suffixed with ".bz2"
 135         output_name = os.path.splitext(input_name)[0]
 136
 137     return output_name
 138
 139 def _do_lzop(input_name, compression=True):
 140     """ Compress/decompress the file with 'lzop' utility.
 141
 142     @input_name: the file name to compress/decompress
 143     @compress: True for compressing, False for decompressing
 144     @retval: the path of the compressed/decompressed file
 145     """
 146     compressor = "lzop"
 147
 148     if compression:
 149         cmdln = [compressor, "-f", "-U", input_name]
 150     else:
 151         cmdln = [compressor, "-d", "-f", "-U", input_name]
 152
 153     _call_external(cmdln)
 154
 155     if compression:
 156         output_name = input_name + ".lzo"
 157     else:
 158         # suppose that file name is suffixed with ".lzo"
 159         output_name = os.path.splitext(input_name)[0]
 160
 161     return output_name
 162
 163 _COMPRESS_SUFFIXES = {
 164     ".lzo"     : [".lzo"],
 165     ".gz"      : [".gz"],
 166     ".bz2"     : [".bz2", ".bz"],
 167     ".tar.lzo" : [".tar.lzo", ".tzo"],
 168     ".tar.gz"  : [".tar.gz", ".tgz", ".taz"],
 169     ".tar.bz2" : [".tar.bz2", ".tbz", ".tbz2", ".tar.bz"],
 170 }
 171
 172 _COMPRESS_FORMATS = {
 173     "gz" :  _do_gzip,
 174     "bz2":  _do_bzip2,
 175     "lzo":  _do_lzop,
 176 }
 177
 178 def get_compress_formats():
 179     """ Get the list of the supported compression formats
 180
 181     @retval: a list contained supported compress formats
 182     """
 183     return _COMPRESS_FORMATS.keys()
 184
 185 def get_compress_suffixes():
 186     """ Get the list of the support suffixes
 187
 188     @retval: a list contained all suffixes
 189     """
 190     suffixes = []
 191     for key in _COMPRESS_SUFFIXES.keys():
 192         suffix = _COMPRESS_SUFFIXES[key]
 193         if (suffix):
 194             suffixes.extend(suffix)
 195
 196     return suffixes
 197
 198 def compress(file_path, compress_format):
 199     """ Compress a given file
 200
 201     @file_path: the path of the file to compress
 202     @compress_format: the compression format
 203     @retval: the path of the compressed file
 204     """
 205     if not os.path.isfile(file_path):
 206         raise OSError, "can't compress a file not existed: '%s'" % file_path
 207
 208     try:
 209         func = _COMPRESS_FORMATS[compress_format]
 210     except KeyError:
 211         raise ValueError, "unknown compress format '%s'" % compress_format
 212     return func(file_path, True)
 213
 214 def decompress(file_path, decompress_format=None):
 215     """ Decompess a give file
 216
 217     @file_path: the path of the file to decompress
 218     @decompress_format: the format for decompression, None for auto detection
 219     @retval: the path of the decompressed file
 220     """
 221     if not os.path.isfile(file_path):
 222         raise OSError, "can't decompress a file not existed: '%s'" % file_path
 223
 224     (file_name, file_ext) = os.path.splitext(file_path)
 225     for key, suffixes in _COMPRESS_SUFFIXES.iteritems():
 226         if file_ext in suffixes:
 227             file_ext = key
 228             break
 229
 230     if file_path != (file_name + file_ext):
 231         shutil.move(file_path, file_name + file_ext)
 232         file_path  = file_name + file_ext
 233
 234     if not decompress_format:
 235         decompress_format = os.path.splitext(file_path)[1].lstrip(".")
 236
 237     try:
 238         func = _COMPRESS_FORMATS[decompress_format]
 239     except KeyError:
 240         raise ValueError, "unknown decompress format '%s'" % decompress_format
 241     return func(file_path, False)
 242
 243
 244 def _do_tar(archive_name, target_name):
 245     """ Archive the directory or the file with 'tar' utility
 246
 247     @archive_name: the name of the tarball file
 248     @target_name: the name of the target to tar
 249     @retval: the path of the archived file
 250     """
 251     file_list = []
 252     if os.path.isdir(target_name):
 253         target_dir = target_name
 254         target_name = "."
 255         for target_file in os.listdir(target_dir):
 256             file_list.append(target_file)
 257     else:
 258         target_dir = os.path.dirname(target_name)
 259         target_name = os.path.basename(target_name)
 260         file_list.append(target_name)
 261     cmdln = ["tar", "-C", target_dir, "-cf", archive_name]
 262     cmdln.extend(file_list)
 263     _call_external(cmdln)
 264
 265     return archive_name
 266
 267 def _do_untar(archive_name, target_dir=None):
 268     """ Unarchive the archived file with 'tar' utility
 269
 270     @archive_name: the name of the archived file
 271     @target_dir: the directory which the target locates
 272     raise exception if fail to extract archive
 273     """
 274     if not target_dir:
 275         target_dir = os.getcwd()
 276     cmdln = ["tar", "-C", target_dir, "-xf", archive_name]
 277     (returncode, stdout, stderr) = _call_external(cmdln)
 278     if returncode != 0:
 279         raise OSError, os.linesep.join([stdout, stderr])
 280
 281 def _imp_tarfile(archive_name, target_name):
 282     """ Archive the directory or the file with tarfile module
 283
 284     @archive_name: the name of the tarball file
 285     @target_name: the name of the target to tar
 286     @retval: the path of the archived file
 287     """
 288     import tarfile
 289
 290     msger.info("Taring files to %s using tarfile module" % archive_name)
 291     tar = tarfile.open(archive_name, 'w')
 292     if os.path.isdir(target_name):
 293         for child in os.listdir(target_name):
 294             tar.add(os.path.join(target_name, child), child)
 295     else:
 296         tar.add(target_name, os.path.basename(target_name))
 297
 298     tar.close()
 299     return archive_name
 300
 301 def _make_tarball(archive_name, target_name, compressor=None):
 302     """ Create a tarball from all the files under 'target_name' or itself.
 303
 304     @archive_name: the name of the archived file to create
 305     @target_name: the directory or the file name to archive
 306     @compressor: callback function to compress the tarball
 307     @retval: indicate the compressing result
 308     """
 309     archive_dir = os.path.dirname(archive_name)
 310     tarball_name = tempfile.mktemp(suffix=".tar", dir=archive_dir)
 311
 312     if which("tar") is not None:
 313         _do_tar(tarball_name, target_name)
 314     else:
 315         _imp_tarfile(tarball_name, target_name)
 316
 317     if compressor:
 318         tarball_name = compressor(tarball_name, True)
 319
 320     shutil.move(tarball_name, archive_name)
 321
 322     return os.path.exists(archive_name)
 323
 324 def _extract_tarball(archive_name, target_dir, compressor=None):
 325     """ Extract a tarball to a target directory
 326
 327     @archive_name: the name of the archived file to extract
 328     @target_dir: the directory of the extracted target
 329     """
 330
 331     _do_untar(archive_name, target_dir)
 332
 333 def _make_zipfile(archive_name, target_name):
 334     """ Create a zip file from all the files under 'target_name' or itself.
 335
 336     @archive_name: the name of the archived file
 337     @target_name: the directory or the file name to archive
 338     @retval: indicate the archiving result
 339     """
 340     import zipfile
 341
 342     msger.info("Zipping files to %s using zipfile module" % archive_name)
 343     arv = zipfile.ZipFile(archive_name, 'w', compression=zipfile.ZIP_DEFLATED)
 344
 345     if os.path.isdir(target_name):
 346         for dirpath, dirname, filenames in os.walk(target_name):
 347             for filename in filenames:
 348                 filepath = os.path.normpath(os.path.join(dirpath, filename))
 349                 arcname = os.path.relpath(filepath, target_name)
 350                 if os.path.isfile(filepath):
 351                     arv.write(filepath, arcname)
 352     else:
 353         arv.write(target_name, os.path.basename(target_name))
 354
 355     arv.close()
 356
 357     return os.path.exists(archive_name)
 358
 359 _ARCHIVE_SUFFIXES = {
 360     "zip"   : [".zip"],
 361     "tar"   : [".tar"],
 362     "lzotar": [".tzo", ".tar.lzo"],
 363     "gztar" : [".tgz", ".taz", ".tar.gz"],
 364     "bztar" : [".tbz", ".tbz2", ".tar.bz", ".tar.bz2"],
 365 }
 366
 367 _ARCHIVE_FORMATS = {
 368     "zip"   : ( _make_zipfile, {} ),
 369     "tar"   : ( _make_tarball, {"compressor" : None} ),
 370     "lzotar": ( _make_tarball, {"compressor" : _do_lzop} ),
 371     "gztar" : ( _make_tarball, {"compressor" : _do_gzip} ),
 372     "bztar" : ( _make_tarball, {"compressor" : _do_bzip2} ),
 373 }
 374
 375 def get_archive_formats():
 376     """ Get the list of the supported formats for archiving
 377
 378     @retval: a list contained archive formats
 379     """
 380     return _ARCHIVE_FORMATS.keys()
 381
 382 def get_archive_suffixes():
 383     """ Get the list of the support suffixes
 384
 385     @retval: a list contained all suffixes
 386     """
 387     suffixes = []
 388     for name in _ARCHIVE_FORMATS.keys():
 389         suffix = _ARCHIVE_SUFFIXES.get(name, None)
 390         if (suffix):
 391             suffixes.extend(suffix)
 392
 393     return suffixes
 394
 395 def make_archive(archive_name, target_name):
 396     """ Create an archive file (eg. tar or zip).
 397
 398     @archive_name: the name of the archived file
 399     @target_name: the directory or the file to archive
 400     @retval: the archiving result
 401     """
 402     if not os.path.exists(target_name):
 403         raise OSError, "archive object does not exist: '%s'" % target_name
 404
 405     for aformat, suffixes in _ARCHIVE_SUFFIXES.iteritems():
 406         if filter(archive_name.endswith, suffixes):
 407             archive_format = aformat
 408             break
 409     else:
 410         raise ValueError, "unknown archive suffix '%s'" % archive_name
 411
 412     try:
 413         func, kwargs = _ARCHIVE_FORMATS[archive_format]
 414     except KeyError:
 415         raise ValueError, "unknown archive format '%s'" % archive_format
 416
 417     archive_name = os.path.abspath(archive_name)
 418     target_name = os.path.abspath(target_name)
 419
 420     archive_dir = os.path.dirname(archive_name)
 421     if not os.path.exists(archive_dir):
 422         os.makedirs(archive_dir)
 423
 424     return func(archive_name, target_name, **kwargs)
 425
 426 def extract_archive(archive_name, target_name):
 427     """ Extract the given file
 428
 429     @archive_name: the name of the archived file to extract
 430     @target_name: the directory name where the target locates
 431     raise exception if fail to extract archive
 432     """
 433     if not os.path.exists(archive_name):
 434         raise OSError, "archived object does not exist: '%s'" % archive_name
 435
 436     archive_name = os.path.abspath(archive_name)
 437     target_name = os.path.abspath(target_name)
 438
 439     if os.path.exists(target_name) and not os.path.isdir(target_name):
 440         raise OSError, "%s should be directory where extracted files locate"\
 441                         % target_name
 442
 443     if not os.path.exists(target_name):
 444         os.makedirs(target_name)
 445
 446     _extract_tarball(archive_name, target_name)
 447
 448 packing = make_archive
 449 compressing = compress
 450