1#!/usr/bin/env python
2
3'''Helper script for constructing an archive (zip or tar) from a list of files.
4
5The output format (tar, tgz, zip) is determined from the file name, unless the user specifies
6--format on the command line.
7
8This script simplifies the specification of filename transformations, so that, e.g.,
9src/mongo/foo.cpp and build/linux2/normal/buildinfo.cpp can get put into the same
10directory in the archive, perhaps mongodb-2.0.2/src/mongo.
11
12Usage:
13
14make_archive.py -o <output-file> [--format (tar|tgz|zip)] \
15    [--transform match1=replacement1 [--transform match2=replacement2 [...]]] \
16    <input file 1> [...]
17
18If the input file names start with "@", the file is expected to contain a list of
19whitespace-separated file names to include in the archive.  This helps get around the Windows
20command line length limit.
21
22Transformations are processed in command-line order and are short-circuiting.  So, if a file matches
23match1, it is never compared against match2 or later.  Matches are just python startswith()
24comparisons.
25
26For a detailed usage example, see src/SConscript.client or src/mongo/SConscript.
27'''
28
29import optparse
30import os
31import sys
32import shlex
33import shutil
34import zipfile
35import tempfile
36from subprocess import (Popen, PIPE, STDOUT)
37
38def main(argv):
39    args = []
40    for arg in argv[1:]:
41        if arg.startswith("@"):
42            file_name = arg[1:]
43            f_handle = open(file_name, "r")
44            args.extend(s1.strip('"') for s1 in shlex.split(f_handle.readline(), posix=False))
45            f_handle.close()
46        else:
47            args.append(arg)
48
49    opts = parse_options(args)
50    if opts.archive_format in ('tar', 'tgz'):
51        make_tar_archive(opts)
52    elif opts.archive_format in ('zip'):
53        make_zip_archive(opts)
54    else:
55        raise ValueError('Unsupported archive format "%s"' % opts.archive_format)
56
57def delete_directory(dir):
58    '''Recursively deletes a directory and its contents.
59    '''
60    try:
61        shutil.rmtree(dir)
62    except Exception:
63        pass
64
65def make_tar_archive(opts):
66    '''Given the parsed options, generates the 'opt.output_filename'
67    tarball containing all the files in 'opt.input_filename' renamed
68    according to the mappings in 'opts.transformations'.
69
70    e.g. for an input file named "a/mongo/build/DISTSRC", and an
71    existing transformation {"a/mongo/build": "release"}, the input
72    file will be written to the tarball as "release/DISTSRC"
73
74    All files to be compressed are copied into new directories as
75    required by 'opts.transformations'. Once the tarball has been
76    created, all temporary directory structures created for the
77    purposes of compressing, are removed.
78    '''
79    tar_options = "cvf"
80    if opts.archive_format is 'tgz':
81        tar_options += "z"
82
83    # clean and create a temp directory to copy files to
84    enclosing_archive_directory = tempfile.mkdtemp(
85        prefix='archive_',
86        dir=os.path.abspath('build')
87    )
88    output_tarfile = os.path.join(os.getcwd(), opts.output_filename)
89
90    tar_command = ["tar", tar_options, output_tarfile]
91
92    for input_filename in opts.input_filenames:
93        preferred_filename = get_preferred_filename(input_filename, opts.transformations)
94        temp_file_location = os.path.join(enclosing_archive_directory, preferred_filename)
95        enclosing_file_directory = os.path.dirname(temp_file_location)
96        if not os.path.exists(enclosing_file_directory):
97            os.makedirs(enclosing_file_directory)
98        print "copying %s => %s" % (input_filename, temp_file_location)
99        if os.path.isdir(input_filename):
100            shutil.copytree(input_filename, temp_file_location)
101        else:
102            shutil.copy2(input_filename, temp_file_location)
103        tar_command.append(preferred_filename)
104
105    print " ".join(tar_command)
106    # execute the full tar command
107    run_directory = os.path.join(os.getcwd(), enclosing_archive_directory)
108    proc = Popen(tar_command, stdout=PIPE, stderr=STDOUT, bufsize=0, cwd=run_directory)
109    proc.wait()
110
111    # delete temp directory
112    delete_directory(enclosing_archive_directory)
113
114def make_zip_archive(opts):
115    '''Given the parsed options, generates the 'opt.output_filename'
116    zipfile containing all the files in 'opt.input_filename' renamed
117    according to the mappings in 'opts.transformations'.
118
119    All files in 'opt.output_filename' are renamed before being
120    written into the zipfile.
121    '''
122    archive = open_zip_archive_for_write(opts.output_filename)
123    try:
124        for input_filename in opts.input_filenames:
125            archive.add(input_filename, arcname=get_preferred_filename(input_filename,
126            opts.transformations))
127    finally:
128        archive.close()
129
130
131def parse_options(args):
132    parser = optparse.OptionParser()
133    parser.add_option('-o', dest='output_filename', default=None,
134                      help='Name of the archive to output.', metavar='FILE')
135    parser.add_option('--format', dest='archive_format', default=None,
136                      choices=('zip', 'tar', 'tgz'),
137                      help='Format of archive to create.  '
138                      'If omitted, use the suffix of the output filename to decide.')
139    parser.add_option('--transform', action='append', dest='transformations', default=[])
140
141    (opts, input_filenames) = parser.parse_args(args)
142    opts.input_filenames = []
143
144    for input_filename in input_filenames:
145        if input_filename.startswith('@'):
146            opts.input_filenames.extend(open(input_filename[1:], 'r').read().split())
147        else:
148            opts.input_filenames.append(input_filename)
149
150    if opts.output_filename is None:
151        parser.error('-o switch is required')
152
153    if opts.archive_format is None:
154        if opts.output_filename.endswith('.zip'):
155            opts.archive_format = 'zip'
156        elif opts.output_filename.endswith('tar.gz') or opts.output_filename.endswith('.tgz'):
157            opts.archive_format = 'tgz'
158        elif opts.output_filename.endswith('.tar'):
159            opts.archive_format = 'tar'
160        else:
161            parser.error('Could not deduce archive format from output filename "%s"' %
162                         opts.output_filename)
163
164    try:
165        opts.transformations = [
166            xform.replace(os.path.altsep or os.path.sep, os.path.sep).split('=', 1)
167            for xform in opts.transformations]
168    except Exception, e:
169        parser.error(e)
170
171    return opts
172
173def open_zip_archive_for_write(filename):
174    '''Open a zip archive for writing and return it.
175    '''
176    # Infuriatingly, Zipfile calls the "add" method "write", but they're otherwise identical,
177    # for our purposes.  WrappedZipFile is a minimal adapter class.
178    class WrappedZipFile(zipfile.ZipFile):
179        def add(self, filename, arcname):
180            return self.write(filename, arcname)
181    return WrappedZipFile(filename, 'w', zipfile.ZIP_DEFLATED)
182
183def get_preferred_filename(input_filename, transformations):
184    '''Does a prefix subsitution on 'input_filename' for the
185    first matching transformation in 'transformations' and
186    returns the substituted string
187    '''
188    for match, replace in transformations:
189        match_lower = match.lower()
190        input_filename_lower = input_filename.lower()
191        if input_filename_lower.startswith(match_lower):
192            return replace + input_filename[len(match):]
193    return input_filename
194
195if __name__ == '__main__':
196    main(sys.argv)
197    sys.exit(0)
198