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            if base_dir != os.curdir:
166                path = os.path.normpath(os.path.join(base_dir, ''))
167                zip.write(path, path)
168                log.info("adding '%s'", path)
169            for dirpath, dirnames, filenames in os.walk(base_dir):
170                for name in dirnames:
171                    path = os.path.normpath(os.path.join(dirpath, name, ''))
172                    zip.write(path, path)
173                    log.info("adding '%s'", path)
174                for name in filenames:
175                    path = os.path.normpath(os.path.join(dirpath, name))
176                    if os.path.isfile(path):
177                        zip.write(path, path)
178                        log.info("adding '%s'" % path)
179            zip.close()
180
181    return zip_filename
182
183ARCHIVE_FORMATS = {
184    'gztar': (make_tarball, [('compress', 'gzip')], "gzip'ed tar-file"),
185    'bztar': (make_tarball, [('compress', 'bzip2')], "bzip2'ed tar-file"),
186    'ztar':  (make_tarball, [('compress', 'compress')], "compressed tar file"),
187    'tar':   (make_tarball, [('compress', None)], "uncompressed tar file"),
188    'zip':   (make_zipfile, [],"ZIP file")
189    }
190
191def check_archive_formats(formats):
192    """Returns the first format from the 'format' list that is unknown.
193
194    If all formats are known, returns None
195    """
196    for format in formats:
197        if format not in ARCHIVE_FORMATS:
198            return format
199    return None
200
201def make_archive(base_name, format, root_dir=None, base_dir=None, verbose=0,
202                 dry_run=0, owner=None, group=None):
203    """Create an archive file (eg. zip or tar).
204
205    'base_name' is the name of the file to create, minus any format-specific
206    extension; 'format' is the archive format: one of "zip", "tar", "ztar",
207    or "gztar".
208
209    'root_dir' is a directory that will be the root directory of the
210    archive; ie. we typically chdir into 'root_dir' before creating the
211    archive.  'base_dir' is the directory where we start archiving from;
212    ie. 'base_dir' will be the common prefix of all files and
213    directories in the archive.  'root_dir' and 'base_dir' both default
214    to the current directory.  Returns the name of the archive file.
215
216    'owner' and 'group' are used when creating a tar archive. By default,
217    uses the current owner and group.
218    """
219    save_cwd = os.getcwd()
220    if root_dir is not None:
221        log.debug("changing into '%s'", root_dir)
222        base_name = os.path.abspath(base_name)
223        if not dry_run:
224            os.chdir(root_dir)
225
226    if base_dir is None:
227        base_dir = os.curdir
228
229    kwargs = {'dry_run': dry_run}
230
231    try:
232        format_info = ARCHIVE_FORMATS[format]
233    except KeyError:
234        raise ValueError, "unknown archive format '%s'" % format
235
236    func = format_info[0]
237    for arg, val in format_info[1]:
238        kwargs[arg] = val
239
240    if format != 'zip':
241        kwargs['owner'] = owner
242        kwargs['group'] = group
243
244    try:
245        filename = func(base_name, base_dir, **kwargs)
246    finally:
247        if root_dir is not None:
248            log.debug("changing back to '%s'", save_cwd)
249            os.chdir(save_cwd)
250
251    return filename
252