~ubuntu-branches/ubuntu/karmic/pypy/karmic

« back to all changes in this revision

Viewing changes to lib-python/2.4.1/zipfile.py

  • Committer: Bazaar Package Importer
  • Author(s): Alexandre Fayolle
  • Date: 2007-04-13 09:33:09 UTC
  • Revision ID: james.westby@ubuntu.com-20070413093309-yoojh4jcoocu2krz
Tags: upstream-1.0.0
ImportĀ upstreamĀ versionĀ 1.0.0

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
"Read and write ZIP files."
 
2
 
 
3
import struct, os, time
 
4
import binascii
 
5
 
 
6
try:
 
7
    import zlib # We may need its compression method
 
8
except ImportError:
 
9
    zlib = None
 
10
 
 
11
__all__ = ["BadZipfile", "error", "ZIP_STORED", "ZIP_DEFLATED", "is_zipfile",
 
12
           "ZipInfo", "ZipFile", "PyZipFile"]
 
13
 
 
14
class BadZipfile(Exception):
 
15
    pass
 
16
error = BadZipfile      # The exception raised by this module
 
17
 
 
18
# constants for Zip file compression methods
 
19
ZIP_STORED = 0
 
20
ZIP_DEFLATED = 8
 
21
# Other ZIP compression methods not supported
 
22
 
 
23
# Here are some struct module formats for reading headers
 
24
structEndArchive = "<4s4H2lH"     # 9 items, end of archive, 22 bytes
 
25
stringEndArchive = "PK\005\006"   # magic number for end of archive record
 
26
structCentralDir = "<4s4B4HlLL5HLl"# 19 items, central directory, 46 bytes
 
27
stringCentralDir = "PK\001\002"   # magic number for central directory
 
28
structFileHeader = "<4s2B4HlLL2H"  # 12 items, file header record, 30 bytes
 
29
stringFileHeader = "PK\003\004"   # magic number for file header
 
30
 
 
31
# indexes of entries in the central directory structure
 
32
_CD_SIGNATURE = 0
 
33
_CD_CREATE_VERSION = 1
 
34
_CD_CREATE_SYSTEM = 2
 
35
_CD_EXTRACT_VERSION = 3
 
36
_CD_EXTRACT_SYSTEM = 4                  # is this meaningful?
 
37
_CD_FLAG_BITS = 5
 
38
_CD_COMPRESS_TYPE = 6
 
39
_CD_TIME = 7
 
40
_CD_DATE = 8
 
41
_CD_CRC = 9
 
42
_CD_COMPRESSED_SIZE = 10
 
43
_CD_UNCOMPRESSED_SIZE = 11
 
44
_CD_FILENAME_LENGTH = 12
 
45
_CD_EXTRA_FIELD_LENGTH = 13
 
46
_CD_COMMENT_LENGTH = 14
 
47
_CD_DISK_NUMBER_START = 15
 
48
_CD_INTERNAL_FILE_ATTRIBUTES = 16
 
49
_CD_EXTERNAL_FILE_ATTRIBUTES = 17
 
50
_CD_LOCAL_HEADER_OFFSET = 18
 
51
 
 
52
# indexes of entries in the local file header structure
 
53
_FH_SIGNATURE = 0
 
54
_FH_EXTRACT_VERSION = 1
 
55
_FH_EXTRACT_SYSTEM = 2                  # is this meaningful?
 
56
_FH_GENERAL_PURPOSE_FLAG_BITS = 3
 
57
_FH_COMPRESSION_METHOD = 4
 
58
_FH_LAST_MOD_TIME = 5
 
59
_FH_LAST_MOD_DATE = 6
 
60
_FH_CRC = 7
 
61
_FH_COMPRESSED_SIZE = 8
 
62
_FH_UNCOMPRESSED_SIZE = 9
 
63
_FH_FILENAME_LENGTH = 10
 
64
_FH_EXTRA_FIELD_LENGTH = 11
 
65
 
 
66
def is_zipfile(filename):
 
67
    """Quickly see if file is a ZIP file by checking the magic number."""
 
68
    try:
 
69
        fpin = open(filename, "rb")
 
70
        endrec = _EndRecData(fpin)
 
71
        fpin.close()
 
72
        if endrec:
 
73
            return True                 # file has correct magic number
 
74
    except IOError:
 
75
        pass
 
76
    return False
 
77
 
 
78
def _EndRecData(fpin):
 
79
    """Return data from the "End of Central Directory" record, or None.
 
80
 
 
81
    The data is a list of the nine items in the ZIP "End of central dir"
 
82
    record followed by a tenth item, the file seek offset of this record."""
 
83
    fpin.seek(-22, 2)               # Assume no archive comment.
 
84
    filesize = fpin.tell() + 22     # Get file size
 
85
    data = fpin.read()
 
86
    if data[0:4] == stringEndArchive and data[-2:] == "\000\000":
 
87
        endrec = struct.unpack(structEndArchive, data)
 
88
        endrec = list(endrec)
 
89
        endrec.append("")               # Append the archive comment
 
90
        endrec.append(filesize - 22)    # Append the record start offset
 
91
        return endrec
 
92
    # Search the last END_BLOCK bytes of the file for the record signature.
 
93
    # The comment is appended to the ZIP file and has a 16 bit length.
 
94
    # So the comment may be up to 64K long.  We limit the search for the
 
95
    # signature to a few Kbytes at the end of the file for efficiency.
 
96
    # also, the signature must not appear in the comment.
 
97
    END_BLOCK = min(filesize, 1024 * 4)
 
98
    fpin.seek(filesize - END_BLOCK, 0)
 
99
    data = fpin.read()
 
100
    start = data.rfind(stringEndArchive)
 
101
    if start >= 0:     # Correct signature string was found
 
102
        endrec = struct.unpack(structEndArchive, data[start:start+22])
 
103
        endrec = list(endrec)
 
104
        comment = data[start+22:]
 
105
        if endrec[7] == len(comment):     # Comment length checks out
 
106
            # Append the archive comment and start offset
 
107
            endrec.append(comment)
 
108
            endrec.append(filesize - END_BLOCK + start)
 
109
            return endrec
 
110
    return      # Error, return None
 
111
 
 
112
 
 
113
class ZipInfo:
 
114
    """Class with attributes describing each file in the ZIP archive."""
 
115
 
 
116
    def __init__(self, filename="NoName", date_time=(1980,1,1,0,0,0)):
 
117
        self.orig_filename = filename   # Original file name in archive
 
118
# Terminate the file name at the first null byte.  Null bytes in file
 
119
# names are used as tricks by viruses in archives.
 
120
        null_byte = filename.find(chr(0))
 
121
        if null_byte >= 0:
 
122
            filename = filename[0:null_byte]
 
123
# This is used to ensure paths in generated ZIP files always use
 
124
# forward slashes as the directory separator, as required by the
 
125
# ZIP format specification.
 
126
        if os.sep != "/":
 
127
            filename = filename.replace(os.sep, "/")
 
128
        self.filename = filename        # Normalized file name
 
129
        self.date_time = date_time      # year, month, day, hour, min, sec
 
130
        # Standard values:
 
131
        self.compress_type = ZIP_STORED # Type of compression for the file
 
132
        self.comment = ""               # Comment for each file
 
133
        self.extra = ""                 # ZIP extra data
 
134
        self.create_system = 0          # System which created ZIP archive
 
135
        self.create_version = 20        # Version which created ZIP archive
 
136
        self.extract_version = 20       # Version needed to extract archive
 
137
        self.reserved = 0               # Must be zero
 
138
        self.flag_bits = 0              # ZIP flag bits
 
139
        self.volume = 0                 # Volume number of file header
 
140
        self.internal_attr = 0          # Internal attributes
 
141
        self.external_attr = 0          # External file attributes
 
142
        # Other attributes are set by class ZipFile:
 
143
        # header_offset         Byte offset to the file header
 
144
        # file_offset           Byte offset to the start of the file data
 
145
        # CRC                   CRC-32 of the uncompressed file
 
146
        # compress_size         Size of the compressed file
 
147
        # file_size             Size of the uncompressed file
 
148
 
 
149
    def FileHeader(self):
 
150
        """Return the per-file header as a string."""
 
151
        dt = self.date_time
 
152
        dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
 
153
        dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2)
 
154
        if self.flag_bits & 0x08:
 
155
            # Set these to zero because we write them after the file data
 
156
            CRC = compress_size = file_size = 0
 
157
        else:
 
158
            CRC = self.CRC
 
159
            compress_size = self.compress_size
 
160
            file_size = self.file_size
 
161
        header = struct.pack(structFileHeader, stringFileHeader,
 
162
                 self.extract_version, self.reserved, self.flag_bits,
 
163
                 self.compress_type, dostime, dosdate, CRC,
 
164
                 compress_size, file_size,
 
165
                 len(self.filename), len(self.extra))
 
166
        return header + self.filename + self.extra
 
167
 
 
168
 
 
169
class ZipFile:
 
170
    """ Class with methods to open, read, write, close, list zip files.
 
171
 
 
172
    z = ZipFile(file, mode="r", compression=ZIP_STORED)
 
173
 
 
174
    file: Either the path to the file, or a file-like object.
 
175
          If it is a path, the file will be opened and closed by ZipFile.
 
176
    mode: The mode can be either read "r", write "w" or append "a".
 
177
    compression: ZIP_STORED (no compression) or ZIP_DEFLATED (requires zlib).
 
178
    """
 
179
 
 
180
    fp = None                   # Set here since __del__ checks it
 
181
 
 
182
    def __init__(self, file, mode="r", compression=ZIP_STORED):
 
183
        """Open the ZIP file with mode read "r", write "w" or append "a"."""
 
184
        if compression == ZIP_STORED:
 
185
            pass
 
186
        elif compression == ZIP_DEFLATED:
 
187
            if not zlib:
 
188
                raise RuntimeError,\
 
189
                      "Compression requires the (missing) zlib module"
 
190
        else:
 
191
            raise RuntimeError, "That compression method is not supported"
 
192
        self.debug = 0  # Level of printing: 0 through 3
 
193
        self.NameToInfo = {}    # Find file info given name
 
194
        self.filelist = []      # List of ZipInfo instances for archive
 
195
        self.compression = compression  # Method of compression
 
196
        self.mode = key = mode.replace('b', '')[0]
 
197
 
 
198
        # Check if we were passed a file-like object
 
199
        if isinstance(file, basestring):
 
200
            self._filePassed = 0
 
201
            self.filename = file
 
202
            modeDict = {'r' : 'rb', 'w': 'wb', 'a' : 'r+b'}
 
203
            self.fp = open(file, modeDict[mode])
 
204
        else:
 
205
            self._filePassed = 1
 
206
            self.fp = file
 
207
            self.filename = getattr(file, 'name', None)
 
208
 
 
209
        if key == 'r':
 
210
            self._GetContents()
 
211
        elif key == 'w':
 
212
            pass
 
213
        elif key == 'a':
 
214
            try:                        # See if file is a zip file
 
215
                self._RealGetContents()
 
216
                # seek to start of directory and overwrite
 
217
                self.fp.seek(self.start_dir, 0)
 
218
            except BadZipfile:          # file is not a zip file, just append
 
219
                self.fp.seek(0, 2)
 
220
        else:
 
221
            if not self._filePassed:
 
222
                self.fp.close()
 
223
                self.fp = None
 
224
            raise RuntimeError, 'Mode must be "r", "w" or "a"'
 
225
 
 
226
    def _GetContents(self):
 
227
        """Read the directory, making sure we close the file if the format
 
228
        is bad."""
 
229
        try:
 
230
            self._RealGetContents()
 
231
        except BadZipfile:
 
232
            if not self._filePassed:
 
233
                self.fp.close()
 
234
                self.fp = None
 
235
            raise
 
236
 
 
237
    def _RealGetContents(self):
 
238
        """Read in the table of contents for the ZIP file."""
 
239
        fp = self.fp
 
240
        endrec = _EndRecData(fp)
 
241
        if not endrec:
 
242
            raise BadZipfile, "File is not a zip file"
 
243
        if self.debug > 1:
 
244
            print endrec
 
245
        size_cd = endrec[5]             # bytes in central directory
 
246
        offset_cd = endrec[6]   # offset of central directory
 
247
        self.comment = endrec[8]        # archive comment
 
248
        # endrec[9] is the offset of the "End of Central Dir" record
 
249
        x = endrec[9] - size_cd
 
250
        # "concat" is zero, unless zip was concatenated to another file
 
251
        concat = x - offset_cd
 
252
        if self.debug > 2:
 
253
            print "given, inferred, offset", offset_cd, x, concat
 
254
        # self.start_dir:  Position of start of central directory
 
255
        self.start_dir = offset_cd + concat
 
256
        fp.seek(self.start_dir, 0)
 
257
        total = 0
 
258
        while total < size_cd:
 
259
            centdir = fp.read(46)
 
260
            total = total + 46
 
261
            if centdir[0:4] != stringCentralDir:
 
262
                raise BadZipfile, "Bad magic number for central directory"
 
263
            centdir = struct.unpack(structCentralDir, centdir)
 
264
            if self.debug > 2:
 
265
                print centdir
 
266
            filename = fp.read(centdir[_CD_FILENAME_LENGTH])
 
267
            # Create ZipInfo instance to store file information
 
268
            x = ZipInfo(filename)
 
269
            x.extra = fp.read(centdir[_CD_EXTRA_FIELD_LENGTH])
 
270
            x.comment = fp.read(centdir[_CD_COMMENT_LENGTH])
 
271
            total = (total + centdir[_CD_FILENAME_LENGTH]
 
272
                     + centdir[_CD_EXTRA_FIELD_LENGTH]
 
273
                     + centdir[_CD_COMMENT_LENGTH])
 
274
            x.header_offset = centdir[_CD_LOCAL_HEADER_OFFSET] + concat
 
275
            # file_offset must be computed below...
 
276
            (x.create_version, x.create_system, x.extract_version, x.reserved,
 
277
                x.flag_bits, x.compress_type, t, d,
 
278
                x.CRC, x.compress_size, x.file_size) = centdir[1:12]
 
279
            x.volume, x.internal_attr, x.external_attr = centdir[15:18]
 
280
            # Convert date/time code to (year, month, day, hour, min, sec)
 
281
            x.date_time = ( (d>>9)+1980, (d>>5)&0xF, d&0x1F,
 
282
                                     t>>11, (t>>5)&0x3F, (t&0x1F) * 2 )
 
283
            self.filelist.append(x)
 
284
            self.NameToInfo[x.filename] = x
 
285
            if self.debug > 2:
 
286
                print "total", total
 
287
        for data in self.filelist:
 
288
            fp.seek(data.header_offset, 0)
 
289
            fheader = fp.read(30)
 
290
            if fheader[0:4] != stringFileHeader:
 
291
                raise BadZipfile, "Bad magic number for file header"
 
292
            fheader = struct.unpack(structFileHeader, fheader)
 
293
            # file_offset is computed here, since the extra field for
 
294
            # the central directory and for the local file header
 
295
            # refer to different fields, and they can have different
 
296
            # lengths
 
297
            data.file_offset = (data.header_offset + 30
 
298
                                + fheader[_FH_FILENAME_LENGTH]
 
299
                                + fheader[_FH_EXTRA_FIELD_LENGTH])
 
300
            fname = fp.read(fheader[_FH_FILENAME_LENGTH])
 
301
            if fname != data.orig_filename:
 
302
                raise RuntimeError, \
 
303
                      'File name in directory "%s" and header "%s" differ.' % (
 
304
                          data.orig_filename, fname)
 
305
 
 
306
    def namelist(self):
 
307
        """Return a list of file names in the archive."""
 
308
        l = []
 
309
        for data in self.filelist:
 
310
            l.append(data.filename)
 
311
        return l
 
312
 
 
313
    def infolist(self):
 
314
        """Return a list of class ZipInfo instances for files in the
 
315
        archive."""
 
316
        return self.filelist
 
317
 
 
318
    def printdir(self):
 
319
        """Print a table of contents for the zip file."""
 
320
        print "%-46s %19s %12s" % ("File Name", "Modified    ", "Size")
 
321
        for zinfo in self.filelist:
 
322
            date = "%d-%02d-%02d %02d:%02d:%02d" % zinfo.date_time
 
323
            print "%-46s %s %12d" % (zinfo.filename, date, zinfo.file_size)
 
324
 
 
325
    def testzip(self):
 
326
        """Read all the files and check the CRC."""
 
327
        for zinfo in self.filelist:
 
328
            try:
 
329
                self.read(zinfo.filename)       # Check CRC-32
 
330
            except BadZipfile:
 
331
                return zinfo.filename
 
332
 
 
333
    def getinfo(self, name):
 
334
        """Return the instance of ZipInfo given 'name'."""
 
335
        return self.NameToInfo[name]
 
336
 
 
337
    def read(self, name):
 
338
        """Return file bytes (as a string) for name."""
 
339
        if self.mode not in ("r", "a"):
 
340
            raise RuntimeError, 'read() requires mode "r" or "a"'
 
341
        if not self.fp:
 
342
            raise RuntimeError, \
 
343
                  "Attempt to read ZIP archive that was already closed"
 
344
        zinfo = self.getinfo(name)
 
345
        filepos = self.fp.tell()
 
346
        self.fp.seek(zinfo.file_offset, 0)
 
347
        bytes = self.fp.read(zinfo.compress_size)
 
348
        self.fp.seek(filepos, 0)
 
349
        if zinfo.compress_type == ZIP_STORED:
 
350
            pass
 
351
        elif zinfo.compress_type == ZIP_DEFLATED:
 
352
            if not zlib:
 
353
                raise RuntimeError, \
 
354
                      "De-compression requires the (missing) zlib module"
 
355
            # zlib compress/decompress code by Jeremy Hylton of CNRI
 
356
            dc = zlib.decompressobj(-15)
 
357
            bytes = dc.decompress(bytes)
 
358
            # need to feed in unused pad byte so that zlib won't choke
 
359
            ex = dc.decompress('Z') + dc.flush()
 
360
            if ex:
 
361
                bytes = bytes + ex
 
362
        else:
 
363
            raise BadZipfile, \
 
364
                  "Unsupported compression method %d for file %s" % \
 
365
            (zinfo.compress_type, name)
 
366
        crc = binascii.crc32(bytes)
 
367
        if crc != zinfo.CRC:
 
368
            raise BadZipfile, "Bad CRC-32 for file %s" % name
 
369
        return bytes
 
370
 
 
371
    def _writecheck(self, zinfo):
 
372
        """Check for errors before writing a file to the archive."""
 
373
        if zinfo.filename in self.NameToInfo:
 
374
            if self.debug:      # Warning for duplicate names
 
375
                print "Duplicate name:", zinfo.filename
 
376
        if self.mode not in ("w", "a"):
 
377
            raise RuntimeError, 'write() requires mode "w" or "a"'
 
378
        if not self.fp:
 
379
            raise RuntimeError, \
 
380
                  "Attempt to write ZIP archive that was already closed"
 
381
        if zinfo.compress_type == ZIP_DEFLATED and not zlib:
 
382
            raise RuntimeError, \
 
383
                  "Compression requires the (missing) zlib module"
 
384
        if zinfo.compress_type not in (ZIP_STORED, ZIP_DEFLATED):
 
385
            raise RuntimeError, \
 
386
                  "That compression method is not supported"
 
387
 
 
388
    def write(self, filename, arcname=None, compress_type=None):
 
389
        """Put the bytes from filename into the archive under the name
 
390
        arcname."""
 
391
        st = os.stat(filename)
 
392
        mtime = time.localtime(st.st_mtime)
 
393
        date_time = mtime[0:6]
 
394
        # Create ZipInfo instance to store file information
 
395
        if arcname is None:
 
396
            zinfo = ZipInfo(filename, date_time)
 
397
        else:
 
398
            zinfo = ZipInfo(arcname, date_time)
 
399
        zinfo.external_attr = (st[0] & 0xFFFF) << 16L      # Unix attributes
 
400
        if compress_type is None:
 
401
            zinfo.compress_type = self.compression
 
402
        else:
 
403
            zinfo.compress_type = compress_type
 
404
        self._writecheck(zinfo)
 
405
        fp = open(filename, "rb")
 
406
        zinfo.flag_bits = 0x00
 
407
        zinfo.header_offset = self.fp.tell()    # Start of header bytes
 
408
        # Must overwrite CRC and sizes with correct data later
 
409
        zinfo.CRC = CRC = 0
 
410
        zinfo.compress_size = compress_size = 0
 
411
        zinfo.file_size = file_size = 0
 
412
        self.fp.write(zinfo.FileHeader())
 
413
        zinfo.file_offset = self.fp.tell()      # Start of file bytes
 
414
        if zinfo.compress_type == ZIP_DEFLATED:
 
415
            cmpr = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
 
416
                 zlib.DEFLATED, -15)
 
417
        else:
 
418
            cmpr = None
 
419
        while 1:
 
420
            buf = fp.read(1024 * 8)
 
421
            if not buf:
 
422
                break
 
423
            file_size = file_size + len(buf)
 
424
            CRC = binascii.crc32(buf, CRC)
 
425
            if cmpr:
 
426
                buf = cmpr.compress(buf)
 
427
                compress_size = compress_size + len(buf)
 
428
            self.fp.write(buf)
 
429
        fp.close()
 
430
        if cmpr:
 
431
            buf = cmpr.flush()
 
432
            compress_size = compress_size + len(buf)
 
433
            self.fp.write(buf)
 
434
            zinfo.compress_size = compress_size
 
435
        else:
 
436
            zinfo.compress_size = file_size
 
437
        zinfo.CRC = CRC
 
438
        zinfo.file_size = file_size
 
439
        # Seek backwards and write CRC and file sizes
 
440
        position = self.fp.tell()       # Preserve current position in file
 
441
        self.fp.seek(zinfo.header_offset + 14, 0)
 
442
        self.fp.write(struct.pack("<lLL", zinfo.CRC, zinfo.compress_size,
 
443
              zinfo.file_size))
 
444
        self.fp.seek(position, 0)
 
445
        self.filelist.append(zinfo)
 
446
        self.NameToInfo[zinfo.filename] = zinfo
 
447
 
 
448
    def writestr(self, zinfo_or_arcname, bytes):
 
449
        """Write a file into the archive.  The contents is the string
 
450
        'bytes'.  'zinfo_or_arcname' is either a ZipInfo instance or
 
451
        the name of the file in the archive."""
 
452
        if not isinstance(zinfo_or_arcname, ZipInfo):
 
453
            zinfo = ZipInfo(filename=zinfo_or_arcname,
 
454
                            date_time=time.localtime(time.time()))
 
455
            zinfo.compress_type = self.compression
 
456
        else:
 
457
            zinfo = zinfo_or_arcname
 
458
        self._writecheck(zinfo)
 
459
        zinfo.file_size = len(bytes)            # Uncompressed size
 
460
        zinfo.CRC = binascii.crc32(bytes)       # CRC-32 checksum
 
461
        if zinfo.compress_type == ZIP_DEFLATED:
 
462
            co = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
 
463
                 zlib.DEFLATED, -15)
 
464
            bytes = co.compress(bytes) + co.flush()
 
465
            zinfo.compress_size = len(bytes)    # Compressed size
 
466
        else:
 
467
            zinfo.compress_size = zinfo.file_size
 
468
        zinfo.header_offset = self.fp.tell()    # Start of header bytes
 
469
        self.fp.write(zinfo.FileHeader())
 
470
        zinfo.file_offset = self.fp.tell()      # Start of file bytes
 
471
        self.fp.write(bytes)
 
472
        if zinfo.flag_bits & 0x08:
 
473
            # Write CRC and file sizes after the file data
 
474
            self.fp.write(struct.pack("<lLL", zinfo.CRC, zinfo.compress_size,
 
475
                  zinfo.file_size))
 
476
        self.filelist.append(zinfo)
 
477
        self.NameToInfo[zinfo.filename] = zinfo
 
478
 
 
479
    def __del__(self):
 
480
        """Call the "close()" method in case the user forgot."""
 
481
        self.close()
 
482
 
 
483
    def close(self):
 
484
        """Close the file, and for mode "w" and "a" write the ending
 
485
        records."""
 
486
        if self.fp is None:
 
487
            return
 
488
        if self.mode in ("w", "a"):             # write ending records
 
489
            count = 0
 
490
            pos1 = self.fp.tell()
 
491
            for zinfo in self.filelist:         # write central directory
 
492
                count = count + 1
 
493
                dt = zinfo.date_time
 
494
                dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
 
495
                dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2)
 
496
                centdir = struct.pack(structCentralDir,
 
497
                  stringCentralDir, zinfo.create_version,
 
498
                  zinfo.create_system, zinfo.extract_version, zinfo.reserved,
 
499
                  zinfo.flag_bits, zinfo.compress_type, dostime, dosdate,
 
500
                  zinfo.CRC, zinfo.compress_size, zinfo.file_size,
 
501
                  len(zinfo.filename), len(zinfo.extra), len(zinfo.comment),
 
502
                  0, zinfo.internal_attr, zinfo.external_attr,
 
503
                  zinfo.header_offset)
 
504
                self.fp.write(centdir)
 
505
                self.fp.write(zinfo.filename)
 
506
                self.fp.write(zinfo.extra)
 
507
                self.fp.write(zinfo.comment)
 
508
            pos2 = self.fp.tell()
 
509
            # Write end-of-zip-archive record
 
510
            endrec = struct.pack(structEndArchive, stringEndArchive,
 
511
                     0, 0, count, count, pos2 - pos1, pos1, 0)
 
512
            self.fp.write(endrec)
 
513
            self.fp.flush()
 
514
        if not self._filePassed:
 
515
            self.fp.close()
 
516
        self.fp = None
 
517
 
 
518
 
 
519
class PyZipFile(ZipFile):
 
520
    """Class to create ZIP archives with Python library files and packages."""
 
521
 
 
522
    def writepy(self, pathname, basename = ""):
 
523
        """Add all files from "pathname" to the ZIP archive.
 
524
 
 
525
        If pathname is a package directory, search the directory and
 
526
        all package subdirectories recursively for all *.py and enter
 
527
        the modules into the archive.  If pathname is a plain
 
528
        directory, listdir *.py and enter all modules.  Else, pathname
 
529
        must be a Python *.py file and the module will be put into the
 
530
        archive.  Added modules are always module.pyo or module.pyc.
 
531
        This method will compile the module.py into module.pyc if
 
532
        necessary.
 
533
        """
 
534
        dir, name = os.path.split(pathname)
 
535
        if os.path.isdir(pathname):
 
536
            initname = os.path.join(pathname, "__init__.py")
 
537
            if os.path.isfile(initname):
 
538
                # This is a package directory, add it
 
539
                if basename:
 
540
                    basename = "%s/%s" % (basename, name)
 
541
                else:
 
542
                    basename = name
 
543
                if self.debug:
 
544
                    print "Adding package in", pathname, "as", basename
 
545
                fname, arcname = self._get_codename(initname[0:-3], basename)
 
546
                if self.debug:
 
547
                    print "Adding", arcname
 
548
                self.write(fname, arcname)
 
549
                dirlist = os.listdir(pathname)
 
550
                dirlist.remove("__init__.py")
 
551
                # Add all *.py files and package subdirectories
 
552
                for filename in dirlist:
 
553
                    path = os.path.join(pathname, filename)
 
554
                    root, ext = os.path.splitext(filename)
 
555
                    if os.path.isdir(path):
 
556
                        if os.path.isfile(os.path.join(path, "__init__.py")):
 
557
                            # This is a package directory, add it
 
558
                            self.writepy(path, basename)  # Recursive call
 
559
                    elif ext == ".py":
 
560
                        fname, arcname = self._get_codename(path[0:-3],
 
561
                                         basename)
 
562
                        if self.debug:
 
563
                            print "Adding", arcname
 
564
                        self.write(fname, arcname)
 
565
            else:
 
566
                # This is NOT a package directory, add its files at top level
 
567
                if self.debug:
 
568
                    print "Adding files from directory", pathname
 
569
                for filename in os.listdir(pathname):
 
570
                    path = os.path.join(pathname, filename)
 
571
                    root, ext = os.path.splitext(filename)
 
572
                    if ext == ".py":
 
573
                        fname, arcname = self._get_codename(path[0:-3],
 
574
                                         basename)
 
575
                        if self.debug:
 
576
                            print "Adding", arcname
 
577
                        self.write(fname, arcname)
 
578
        else:
 
579
            if pathname[-3:] != ".py":
 
580
                raise RuntimeError, \
 
581
                      'Files added with writepy() must end with ".py"'
 
582
            fname, arcname = self._get_codename(pathname[0:-3], basename)
 
583
            if self.debug:
 
584
                print "Adding file", arcname
 
585
            self.write(fname, arcname)
 
586
 
 
587
    def _get_codename(self, pathname, basename):
 
588
        """Return (filename, archivename) for the path.
 
589
 
 
590
        Given a module name path, return the correct file path and
 
591
        archive name, compiling if necessary.  For example, given
 
592
        /python/lib/string, return (/python/lib/string.pyc, string).
 
593
        """
 
594
        file_py  = pathname + ".py"
 
595
        file_pyc = pathname + ".pyc"
 
596
        file_pyo = pathname + ".pyo"
 
597
        if os.path.isfile(file_pyo) and \
 
598
                            os.stat(file_pyo).st_mtime >= os.stat(file_py).st_mtime:
 
599
            fname = file_pyo    # Use .pyo file
 
600
        elif not os.path.isfile(file_pyc) or \
 
601
             os.stat(file_pyc).st_mtime < os.stat(file_py).st_mtime:
 
602
            import py_compile
 
603
            if self.debug:
 
604
                print "Compiling", file_py
 
605
            try:
 
606
                py_compile.compile(file_py, file_pyc, None, True)
 
607
            except py_compile.PyCompileError,err:
 
608
                print err.msg
 
609
            fname = file_pyc
 
610
        else:
 
611
            fname = file_pyc
 
612
        archivename = os.path.split(fname)[1]
 
613
        if basename:
 
614
            archivename = "%s/%s" % (basename, archivename)
 
615
        return (fname, archivename)