1"""distutils.archive_util
2
3Utility functions for creating archive files (tarballs, zip files,
4that sort of thing)."""
5
6__revision__ = "$Id$"
7
8import os
9from warnings import warn
10import sys
11
12from distutils.errors import DistutilsExecError
13from distutils.spawn import spawn
14from distutils.dir_util import mkpath
15from distutils import log
16
17try:
18    from pwd import getpwnam
19except ImportError:
20    getpwnam = None
21
22try:
23    from grp import getgrnam
24except ImportError:
25    getgrnam = None
26
27def _get_gid(name):
28    """Returns a gid, given a group name."""
29    if getgrnam is None or name is None:
30        return None
31    try:
32        result = getgrnam(name)
33    except KeyError:
34        result = None
35    if result is not None:
36        return result[2]
37    return None
38
39def _get_uid(name):
40    """Returns an uid, given a user name."""
41    if getpwnam is None or name is None:
42        return None
43    try:
44        result = getpwnam(name)
45    except KeyError:
46        result = None
47    if result is not None:
48        return result[2]
49    return None
50
51def make_tarball(base_name, base_dir, compress="gzip", verbose=0, dry_run=0,
52                 owner=None, group=None):
53    """Create a (possibly compressed) tar file from all the files under
54    'base_dir'.
55
56    'compress' must be "gzip" (the default), "compress", "bzip2", or None.
57    (compress will be deprecated in Python 3.2)
58
59    'owner' and 'group' can be used to define an owner and a group for the
60    archive that is being built. If not provided, the current owner and group
61    will be used.
62
63    The output tar file will be named 'base_dir' +  ".tar", possibly plus
64    the appropriate compression extension (".gz", ".bz2" or ".Z").
65
66    Returns the output filename.
67    """
68    tar_compression = {'gzip': 'gz', 'bzip2': 'bz2', None: '', 'compress': ''}
69    compress_ext = {'gzip': '.gz', 'bzip2': '.bz2', 'compress': '.Z'}
70
71    # flags for compression program, each element of list will be an argument
72    if compress is not None and compress not in compress_ext.keys():
73        raise ValueError, \
74              ("bad value for 'compress': must be None, 'gzip', 'bzip2' "
75               "or 'compress'")
76
77    archive_name = base_name + '.tar'
78    if compress != 'compress':
79        archive_name += compress_ext.get(compress, '')
80
81    mkpath(os.path.dirname(archive_name), dry_run=dry_run)
82
83    # creating the tarball
84    import tarfile  # late import so Python build itself doesn't break
85
86    log.info('Creating tar archive')
87
88    uid = _get_uid(owner)
89    gid = _get_gid(group)
90
91    def _set_uid_gid(tarinfo):
92        if gid is not None:
93            tarinfo.gid = gid
94            tarinfo.gname = group
95        if uid is not None:
96            tarinfo.uid = uid
97            tarinfo.uname = owner
98        return tarinfo
99
100    if not dry_run:
101        tar = tarfile.open(archive_name, 'w|%s' % tar_compression[compress])
102        try:
103            tar.add(base_dir, filter=_set_uid_gid)
104        finally:
105            tar.close()
106
107    # compression using `compress`
108    if compress == 'compress':
109        warn("'compress' will be deprecated.", PendingDeprecationWarning)
110        # the option varies depending on the platform
111        compressed_name = archive_name + compress_ext[compress]
112        if sys.platform == 'win32':
113            cmd = [compress, archive_name, compressed_name]
114        else:
115            cmd = [compress, '-f', archive_name]
116        spawn(cmd, dry_run=dry_run)
117        return compressed_name
118
119    return archive_name
120
121def make_zipfile(base_name, base_dir, verbose=0, dry_run=0):
122    """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    "zipfile" Python module (if available) or the InfoZIP "zip" utility
126    (if installed and found on the default search path).  If neither tool is
127    available, raises DistutilsExecError.  Returns the name of the output zip
128    file.
129    """
130    try:
131        import zipfile
132    except ImportError:
133        zipfile = None
134
135    zip_filename = base_name + ".zip"
136    mkpath(os.path.dirname(zip_filename), dry_run=dry_run)
137
138    # If zipfile module is not available, try spawning an external
139    # 'zip' command.
140    if zipfile is None:
141        if verbose:
142            zipoptions = "-r"
143        else:
144            zipoptions = "-rq"
145
146        try:
147            spawn(["zip", zipoptions, zip_filename, base_dir],
148                  dry_run=dry_run)
149        except DistutilsExecError:
150            # XXX really should distinguish between "couldn't find
151            # external 'zip' command" and "zip failed".
152            raise DistutilsExecError, \
153                  ("unable to create zip file '%s': "
154                   "could neither import the 'zipfile' module nor "
155                   "find a standalone zip utility") % zip_filename
156
157    else:
158        log.info("creating '%s' and adding '%s' to it",
159                 zip_filename, base_dir)
160
161        if not dry_run:
162            zip = zipfile.ZipFile(zip_filename, "w",
163                                  compression=zipfile.ZIP_DEFLATED)
164
165            for dirpath, dirnames, filenames in os.walk(base_dir):
166                for name in filenames:
167                    path = os.path.normpath(os.path.join(dirpath, name))
168                    if os.path.isfile(path):
169                        zip.write(path, path)
170                        log.info("adding '%s'" % path)
171            zip.close()
172
173    return zip_filename
174
175ARCHIVE_FORMATS = {
176    'gztar': (make_tarball, [('compress', 'gzip')], "gzip'ed tar-file"),
177    'bztar': (make_tarball, [('compress', 'bzip2')], "bzip2'ed tar-file"),
178    'ztar':  (make_tarball, [('compress', 'compress')], "compressed tar file"),
179    'tar':   (make_tarball, [('compress', None)], "uncompressed tar file"),
180    'zip':   (make_zipfile, [],"ZIP file")
181    }
182
183def check_archive_formats(formats):
184    """Returns the first format from the 'format' list that is unknown.
185
186    If all formats are known, returns None
187    """
188    for format in formats:
189        if format not in ARCHIVE_FORMATS:
190            return format
191    return None
192
193def make_archive(base_name, format, root_dir=None, base_dir=None, verbose=0,
194                 dry_run=0, owner=None, group=None):
195    """Create an archive file (eg. zip or tar).
196
197    'base_name' is the name of the file to create, minus any format-specific
198    extension; 'format' is the archive format: one of "zip", "tar", "ztar",
199    or "gztar".
200
201    'root_dir' is a directory that will be the root directory of the
202    archive; ie. we typically chdir into 'root_dir' before creating the
203    archive.  'base_dir' is the directory where we start archiving from;
204    ie. 'base_dir' will be the common prefix of all files and
205    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    'owner' and 'group' are used when creating a tar archive. By default,
209    uses the current owner and group.
210    """
211    save_cwd = os.getcwd()
212    if root_dir is not None:
213        log.debug("changing into '%s'", root_dir)
214        base_name = os.path.abspath(base_name)
215        if not dry_run:
216            os.chdir(root_dir)
217
218    if base_dir is None:
219        base_dir = os.curdir
220
221    kwargs = {'dry_run': dry_run}
222
223    try:
224        format_info = ARCHIVE_FORMATS[format]
225    except KeyError:
226        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    if format != 'zip':
233        kwargs['owner'] = owner
234        kwargs['group'] = group
235
236    try:
237        filename = func(base_name, base_dir, **kwargs)
238    finally:
239        if root_dir is not None:
240            log.debug("changing back to '%s'", save_cwd)
241            os.chdir(save_cwd)
242
243    return filename
244