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