archive_util.py 7.64 KB
Newer Older
1 2 3 4 5 6 7 8
"""distutils.archive_util

Utility functions for creating archive files (tarballs, zip files,
that sort of thing)."""

__revision__ = "$Id$"

import os
9 10 11
from warnings import warn
import sys

12 13
from distutils.errors import DistutilsExecError
from distutils.spawn import spawn
14
from distutils.dir_util import mkpath
15
from distutils import log
16

17 18
try:
    from pwd import getpwnam
19
except ImportError:
20 21 22 23
    getpwnam = None

try:
    from grp import getgrnam
24
except ImportError:
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
    getgrnam = None

def _get_gid(name):
    """Returns a gid, given a group name."""
    if getgrnam is None or name is None:
        return None
    try:
        result = getgrnam(name)
    except KeyError:
        result = None
    if result is not None:
        return result[2]
    return None

def _get_uid(name):
    """Returns an uid, given a user name."""
    if getpwnam is None or name is None:
        return None
    try:
        result = getpwnam(name)
    except KeyError:
        result = None
    if result is not None:
        return result[2]
    return None

def make_tarball(base_name, base_dir, compress="gzip", verbose=0, dry_run=0,
                 owner=None, group=None):
53
    """Create a (possibly compressed) tar file from all the files under
54 55 56
    'base_dir'.

    'compress' must be "gzip" (the default), "compress", "bzip2", or None.
57 58 59 60 61 62
    (compress will be deprecated in Python 3.2)

    'owner' and 'group' can be used to define an owner and a group for the
    archive that is being built. If not provided, the current owner and group
    will be used.

63 64
    The output tar file will be named 'base_dir' +  ".tar", possibly plus
    the appropriate compression extension (".gz", ".bz2" or ".Z").
65

66
    Returns the output filename.
Greg Ward's avatar
Greg Ward committed
67
    """
68 69
    tar_compression = {'gzip': 'gz', 'bzip2': 'bz2', None: '', 'compress': ''}
    compress_ext = {'gzip': '.gz', 'bzip2': '.bz2', 'compress': '.Z'}
Fred Drake's avatar
Fred Drake committed
70

Greg Ward's avatar
Greg Ward committed
71 72
    # flags for compression program, each element of list will be an argument
    if compress is not None and compress not in compress_ext.keys():
73
        raise ValueError, \
74 75 76 77 78 79
              ("bad value for 'compress': must be None, 'gzip', 'bzip2' "
               "or 'compress'")

    archive_name = base_name + '.tar'
    if compress != 'compress':
        archive_name += compress_ext.get(compress, '')
80

81
    mkpath(os.path.dirname(archive_name), dry_run=dry_run)
82 83 84 85 86

    # creating the tarball
    import tarfile  # late import so Python build itself doesn't break

    log.info('Creating tar archive')
87 88 89 90 91 92 93 94 95 96 97 98 99

    uid = _get_uid(owner)
    gid = _get_gid(group)

    def _set_uid_gid(tarinfo):
        if gid is not None:
            tarinfo.gid = gid
            tarinfo.gname = group
        if uid is not None:
            tarinfo.uid = uid
            tarinfo.uname = owner
        return tarinfo

100 101 102
    if not dry_run:
        tar = tarfile.open(archive_name, 'w|%s' % tar_compression[compress])
        try:
103
            tar.add(base_dir, filter=_set_uid_gid)
104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119
        finally:
            tar.close()

    # compression using `compress`
    if compress == 'compress':
        warn("'compress' will be deprecated.", PendingDeprecationWarning)
        # the option varies depending on the platform
        compressed_name = archive_name + compress_ext[compress]
        if sys.platform == 'win32':
            cmd = [compress, archive_name, compressed_name]
        else:
            cmd = [compress, '-f', archive_name]
        spawn(cmd, dry_run=dry_run)
        return compressed_name

    return archive_name
120

121 122
def make_zipfile(base_name, base_dir, verbose=0, dry_run=0):
    """Create a zip file from all the files under 'base_dir'.
123

124
    The output zip file will be named 'base_name' + ".zip".  Uses either the
125 126 127 128
    "zipfile" Python module (if available) or the InfoZIP "zip" utility
    (if installed and found on the default search path).  If neither tool is
    available, raises DistutilsExecError.  Returns the name of the output zip
    file.
Greg Ward's avatar
Greg Ward committed
129
    """
130 131 132 133
    try:
        import zipfile
    except ImportError:
        zipfile = None
134

135
    zip_filename = base_name + ".zip"
136
    mkpath(os.path.dirname(zip_filename), dry_run=dry_run)
137 138 139 140 141 142 143 144

    # If zipfile module is not available, try spawning an external
    # 'zip' command.
    if zipfile is None:
        if verbose:
            zipoptions = "-r"
        else:
            zipoptions = "-rq"
145

146
        try:
147 148 149 150 151
            spawn(["zip", zipoptions, zip_filename, base_dir],
                  dry_run=dry_run)
        except DistutilsExecError:
            # XXX really should distinguish between "couldn't find
            # external 'zip' command" and "zip failed".
152
            raise DistutilsExecError, \
153 154 155
                  ("unable to create zip file '%s': "
                   "could neither import the 'zipfile' module nor "
                   "find a standalone zip utility") % zip_filename
156

157 158
    else:
        log.info("creating '%s' and adding '%s' to it",
159
                 zip_filename, base_dir)
160

161
        if not dry_run:
162 163
            zip = zipfile.ZipFile(zip_filename, "w",
                                  compression=zipfile.ZIP_DEFLATED)
164

165 166 167 168
            for dirpath, dirnames, filenames in os.walk(base_dir):
                for name in filenames:
                    path = os.path.normpath(os.path.join(dirpath, name))
                    if os.path.isfile(path):
169
                        zip.write(path, path)
170
                        log.info("adding '%s'" % path)
171
            zip.close()
172 173 174

    return zip_filename

175
ARCHIVE_FORMATS = {
176 177 178 179
    'gztar': (make_tarball, [('compress', 'gzip')], "gzip'ed tar-file"),
    'bztar': (make_tarball, [('compress', 'bzip2')], "bzip2'ed tar-file"),
    'ztar':  (make_tarball, [('compress', 'compress')], "compressed tar file"),
    'tar':   (make_tarball, [('compress', None)], "uncompressed tar file"),
180
    'zip':   (make_zipfile, [],"ZIP file")
181 182
    }

183 184 185 186 187
def check_archive_formats(formats):
    """Returns the first format from the 'format' list that is unknown.

    If all formats are known, returns None
    """
188
    for format in formats:
189
        if format not in ARCHIVE_FORMATS:
190
            return format
191 192 193
    return None

def make_archive(base_name, format, root_dir=None, base_dir=None, verbose=0,
194
                 dry_run=0, owner=None, group=None):
195 196 197 198 199 200
    """Create an archive file (eg. zip or tar).

    'base_name' is the name of the file to create, minus any format-specific
    extension; 'format' is the archive format: one of "zip", "tar", "ztar",
    or "gztar".

201 202 203 204 205
    'root_dir' is a directory that will be the root directory of the
    archive; ie. we typically chdir into 'root_dir' before creating the
    archive.  'base_dir' is the directory where we start archiving from;
    ie. 'base_dir' will be the common prefix of all files and
    directories in the archive.  'root_dir' and 'base_dir' both default
206
    to the current directory.  Returns the name of the archive file.
207 208 209

    'owner' and 'group' are used when creating a tar archive. By default,
    uses the current owner and group.
210
    """
211 212
    save_cwd = os.getcwd()
    if root_dir is not None:
213
        log.debug("changing into '%s'", root_dir)
Greg Ward's avatar
Greg Ward committed
214
        base_name = os.path.abspath(base_name)
215
        if not dry_run:
Greg Ward's avatar
Greg Ward committed
216
            os.chdir(root_dir)
217 218 219 220

    if base_dir is None:
        base_dir = os.curdir

221
    kwargs = {'dry_run': dry_run}
Fred Drake's avatar
Fred Drake committed
222

223 224 225 226
    try:
        format_info = ARCHIVE_FORMATS[format]
    except KeyError:
        raise ValueError, "unknown archive format '%s'" % format
227

228
    func = format_info[0]
229
    for arg, val in format_info[1]:
230
        kwargs[arg] = val
231

232 233 234 235
    if format != 'zip':
        kwargs['owner'] = owner
        kwargs['group'] = group

236 237 238 239 240 241
    try:
        filename = func(base_name, base_dir, **kwargs)
    finally:
        if root_dir is not None:
            log.debug("changing back to '%s'", save_cwd)
            os.chdir(save_cwd)
242

243
    return filename