ZipPythonDir.py 7.33 KB
Newer Older
marcocle's avatar
marcocle committed
1
2
#!/usr/bin/env python

Marco Clemencic's avatar
Marco Clemencic committed
3
# file ZipPythonDir.py
marcocle's avatar
marcocle committed
4
#  Script to generate a zip file that can replace a directory in the python path.
5

6
7
8
9
10
11
12
13
import os
import sys
import zipfile
import logging
import stat
import time
import re
import codecs
marcocle's avatar
marcocle committed
14
15
from StringIO import StringIO

Marco Clemencic's avatar
Marco Clemencic committed
16
17
18
# Class for generic exception coming from the zipdir() function


marcocle's avatar
marcocle committed
19
20
21
class ZipdirError(RuntimeError):
    pass

Gitlab CI's avatar
Gitlab CI committed
22

Marco Clemencic's avatar
Marco Clemencic committed
23
# Collect the changes to be applied to the zip file.
marcocle's avatar
marcocle committed
24
25
26
27
28
#
#  @param directory: directory to be packed in the zip file
#  @param infolist: list of ZipInfo objects already contained in the zip archive
#
#  @return: tuple of (added, modified, untouched, removed) entries in the directory with respect to the zip file
29
#
Marco Clemencic's avatar
Marco Clemencic committed
30
31


marcocle's avatar
marcocle committed
32
33
34
35
36
37
38
39
def _zipChanges(directory, infolist):
    # gets the dates of the files in the zip archive
    infos = {}
    for i in infolist:
        fn = i.filename
        if fn.endswith(".pyc"):
            fn = fn[:-1]
        infos[fn] = i.date_time
40

marcocle's avatar
marcocle committed
41
42
43
44
45
46
    # gets the changes
    added = []
    modified = []
    untouched = []
    removed = []
    all_files = set()
47

marcocle's avatar
marcocle committed
48
49
    log = logging.getLogger("zipdir")
    dirlen = len(directory) + 1
50
51
52
53
    for root, dirs, files in os.walk(directory):
        if "lib-dynload" in dirs:
            # exclude the directory containing binary modules
            dirs.remove("lib-dynload")
marcocle's avatar
marcocle committed
54
55
56
        arcdir = root[dirlen:]
        for f in files:
            ext = os.path.splitext(f)[1]
Marco Clemencic's avatar
Marco Clemencic committed
57
            if ext == ".py":  # extensions that can enter the zip file
marcocle's avatar
marcocle committed
58
59
60
61
62
63
                filename = os.path.join(arcdir, f)
                all_files.add(filename)
                if filename not in infos:
                    action = "A"
                    added.append(filename)
                else:
Marco Clemencic's avatar
Marco Clemencic committed
64
                    filetime = time.localtime(
Gitlab CI's avatar
Gitlab CI committed
65
66
                        os.stat(os.path.join(directory,
                                             filename))[stat.ST_MTIME])[:6]
marcocle's avatar
marcocle committed
67
68
69
70
71
72
                    if filetime > infos[filename]:
                        action = "M"
                        modified.append(filename)
                    else:
                        action = "U"
                        untouched.append(filename)
marcocle's avatar
marcocle committed
73
74
75
76
                if action in ['U']:
                    log.debug(" %s -> %s", action, filename)
                else:
                    log.info(" %s -> %s", action, filename)
77
            # cases that can be ignored
78
79
            elif (ext not in [".pyc", ".pyo", ".stamp", ".cmtref", ".confdb"]
                  and not f.startswith('.__afs')):
Marco Clemencic's avatar
Marco Clemencic committed
80
                raise ZipdirError(
Gitlab CI's avatar
Gitlab CI committed
81
82
                    "Cannot add '%s' to the zip file, only '.py' are allowed."
                    % os.path.join(arcdir, f))
marcocle's avatar
marcocle committed
83
84
85
86
87
88
89
    # check for removed files
    for filename in infos:
        if filename not in all_files:
            removed.append(filename)
            log.info(" %s -> %s", "R", filename)
    return (added, modified, untouched, removed)

Marco Clemencic's avatar
Marco Clemencic committed
90

91
def checkEncoding(fileObj):
92
93
94
95
    '''
    Check that a file honors the declared encoding (default ASCII for Python 2
    and UTF-8 for Python 3).

96
97
    Raises a UnicodeDecodeError in case of decoding problems and LookupError if
    the specified codec does not exists.
98
99
100

    See http://www.python.org/dev/peps/pep-0263/
    '''
101
102
    from itertools import islice

103
104
105
106
107
108
    # default encoding
    if sys.version_info[0] <= 2:
        enc = 'ascii'
    else:
        enc = 'utf-8'

109
    # find the encoding of the file, if specified (in the first two lines)
110
    enc_exp = re.compile(r"coding[:=]\s*([-\w.]+)")
111
112
    for l in islice(fileObj, 2):
        m = enc_exp.search(l)
113
114
115
116
        if m:
            enc = m.group(1)
            break

117
118
119
120
    if hasattr(fileObj, 'name'):
        logging.getLogger('checkEncoding').debug('checking encoding %s on %s',
                                                 enc, fileObj.name)
    else:
Gitlab CI's avatar
Gitlab CI committed
121
122
        logging.getLogger('checkEncoding').debug(
            'checking encoding %s on file object', enc)
123
    # try to read the file with the declared encoding
124
125
    fileObj.seek(0)
    codecs.getreader(enc)(fileObj).read()
126
127


Marco Clemencic's avatar
Marco Clemencic committed
128
129
# Make a zip file out of a directory containing python modules
def zipdir(directory, no_pyc=False):
Marco Clemencic's avatar
Marco Clemencic committed
130
    filename = os.path.realpath(directory + ".zip")
marcocle's avatar
marcocle committed
131
132
    log = logging.getLogger("zipdir")
    if not os.path.isdir(directory):
Gitlab CI's avatar
Gitlab CI committed
133
134
        log.warning('directory %s missing, creating empty .zip file',
                    directory)
Marco Clemencic's avatar
Marco Clemencic committed
135
136
        open(filename, "ab").close()
        return
marcocle's avatar
marcocle committed
137
138
139
140
    msg = "Zipping directory '%s'"
    if no_pyc:
        msg += " (without pre-compilation)"
    log.info(msg, directory)
141

marcocle's avatar
marcocle committed
142
143
144
145
146
147
148
149
    # Open the file in read an update mode
    if os.path.exists(filename):
        zipFile = open(filename, "r+b")
    else:
        # If the file does not exist, we need to create it.
        # "append mode" ensures that, in case of two processes trying to
        # create the file, they do not truncate each other file
        zipFile = open(filename, "ab")
150

marcocle's avatar
marcocle committed
151
152
153
154
155
    try:
        if zipfile.is_zipfile(filename):
            infolist = zipfile.ZipFile(filename).infolist()
        else:
            infolist = []
Gitlab CI's avatar
Gitlab CI committed
156
157
        (added, modified, untouched, removed) = _zipChanges(
            directory, infolist)
marcocle's avatar
marcocle committed
158
159
160
161
        if added or modified or removed:
            tempBuf = StringIO()
            z = zipfile.PyZipFile(tempBuf, "w", zipfile.ZIP_DEFLATED)
            for f in added + modified + untouched:
marcocle's avatar
marcocle committed
162
                src = os.path.join(directory, f)
163
                checkEncoding(open(src, 'rb'))
marcocle's avatar
marcocle committed
164
                if no_pyc:
marcocle's avatar
marcocle committed
165
166
                    log.debug("adding '%s'", f)
                    z.write(src, f)
marcocle's avatar
marcocle committed
167
                else:
marcocle's avatar
marcocle committed
168
169
170
171
172
173
                    # Remove the .pyc file to always force a re-compilation
                    if os.path.exists(src + 'c'):
                        log.debug("removing old .pyc for '%s'", f)
                        os.remove(src + 'c')
                    log.debug("adding '%s'", f)
                    z.writepy(src, os.path.dirname(f))
marcocle's avatar
marcocle committed
174
175
176
177
178
179
180
            z.close()
            zipFile.seek(0)
            zipFile.write(tempBuf.getvalue())
            zipFile.truncate()
            log.info("File '%s' closed", filename)
        else:
            log.info("Nothing to do on '%s'", filename)
181
    except UnicodeDecodeError as x:
182
183
184
185
        log.error("Wrong encoding in file '%s':", src)
        log.error("    %s", x)
        log.error("Probably you forgot the line '# -*- coding: utf-8 -*-'")
        sys.exit(1)
marcocle's avatar
marcocle committed
186
187
188
    finally:
        zipFile.close()

Gitlab CI's avatar
Gitlab CI committed
189

Marco Clemencic's avatar
Marco Clemencic committed
190
# Main function of the script.
marcocle's avatar
marcocle committed
191
#  Parse arguments and call zipdir() for each directory passed as argument
Marco Clemencic's avatar
Marco Clemencic committed
192
193
194


def main(argv=None):
marcocle's avatar
marcocle committed
195
    from optparse import OptionParser
Marco Clemencic's avatar
Marco Clemencic committed
196
    parser = OptionParser(usage="%prog [options] directory1 [directory2 ...]")
Gitlab CI's avatar
Gitlab CI committed
197
198
199
200
201
202
203
204
205
206
    parser.add_option(
        "--no-pyc",
        action="store_true",
        help="copy the .py files without pre-compiling them")
    parser.add_option(
        "--quiet", action="store_true", help="do not print info messages")
    parser.add_option(
        "--debug",
        action="store_true",
        help="print debug messages (has priority over --quiet)")
207

marcocle's avatar
marcocle committed
208
209
210
    if argv is None:
        argv = sys.argv
    opts, args = parser.parse_args(argv[1:])
211

marcocle's avatar
marcocle committed
212
213
    if not args:
        parser.error("Specify at least one directory to zip")
214

marcocle's avatar
marcocle committed
215
216
217
218
    # Initialize the logging module
    level = logging.INFO
    if opts.quiet:
        level = logging.WARNING
marcocle's avatar
marcocle committed
219
220
    if opts.debug:
        level = logging.DEBUG
221
    logging.basicConfig(level=level)
222

marcocle's avatar
marcocle committed
223
224
225
226
    # zip all the directories passed as arguments
    for d in args:
        zipdir(d, opts.no_pyc)

Marco Clemencic's avatar
Marco Clemencic committed
227

marcocle's avatar
marcocle committed
228
229
if __name__ == '__main__':
    main()