~landscape/zope3/newer-from-ztk

« back to all changes in this revision

Viewing changes to src/twisted/python/zipstream.py

  • Committer: Thomas Hervé
  • Date: 2009-07-08 13:52:04 UTC
  • Revision ID: thomas@canonical.com-20090708135204-df5eesrthifpylf8
Remove twisted copy

Show diffs side-by-side

added added

removed removed

Lines of Context:
1
 
"""An extremely asynch approach to unzipping files.  This allows you
2
 
to unzip a little bit of a file at a time, which means it can
3
 
integrate nicely with a reactor.
4
 
 
5
 
"""
6
 
 
7
 
from __future__ import generators
8
 
 
9
 
import zipfile
10
 
import os.path
11
 
import binascii
12
 
import zlib
13
 
import struct
14
 
 
15
 
class ChunkingZipFile(zipfile.ZipFile):
16
 
    """A ZipFile object which, with readfile(), also gives you access
17
 
    to a filelike object for each entry.
18
 
    """
19
 
    def readfile(self, name):
20
 
        """Return file-like object for name."""
21
 
        if self.mode not in ("r", "a"):
22
 
            raise RuntimeError, 'read() requires mode "r" or "a"'
23
 
        if not self.fp:
24
 
            raise RuntimeError, \
25
 
                  "Attempt to read ZIP archive that was already closed"
26
 
        zinfo = self.getinfo(name)
27
 
 
28
 
        self.fp.seek(zinfo.header_offset, 0)
29
 
 
30
 
        # Skip the file header:
31
 
        fheader = self.fp.read(30)
32
 
        if fheader[0:4] != zipfile.stringFileHeader:
33
 
            raise zipfile.BadZipfile, "Bad magic number for file header"
34
 
 
35
 
        fheader = struct.unpack(zipfile.structFileHeader, fheader)
36
 
        fname = self.fp.read(fheader[zipfile._FH_FILENAME_LENGTH])
37
 
        if fheader[zipfile._FH_EXTRA_FIELD_LENGTH]:
38
 
            self.fp.read(fheader[zipfile._FH_EXTRA_FIELD_LENGTH])
39
 
 
40
 
        if fname != zinfo.orig_filename:
41
 
            raise zipfile.BadZipfile, \
42
 
                      'File name in directory "%s" and header "%s" differ.' % (
43
 
                          zinfo.orig_filename, fname)
44
 
 
45
 
        if zinfo.compress_type == zipfile.ZIP_STORED:
46
 
            return ZipFileEntry(self.fp, zinfo.compress_size)
47
 
        elif zinfo.compress_type == zipfile.ZIP_DEFLATED:
48
 
            if not zlib:
49
 
                raise RuntimeError, \
50
 
                      "De-compression requires the (missing) zlib module"
51
 
            return DeflatedZipFileEntry(self.fp, zinfo.compress_size)
52
 
        else:
53
 
            raise zipfile.BadZipfile, \
54
 
                  "Unsupported compression method %d for file %s" % \
55
 
            (zinfo.compress_type, name)
56
 
    
57
 
    def read(self, name):
58
 
        """Return file bytes (as a string) for name."""
59
 
        f = self.readfile(name)
60
 
        zinfo = self.getinfo(name)
61
 
        bytes = f.read()
62
 
        crc = binascii.crc32(bytes)
63
 
        if crc != zinfo.CRC:
64
 
            raise zipfile.BadZipfile, "Bad CRC-32 for file %s" % name
65
 
        return bytes        
66
 
 
67
 
 
68
 
class ZipFileEntry:
69
 
    """File-like object used to read an uncompressed entry in a ZipFile"""
70
 
    
71
 
    def __init__(self, fp, length):
72
 
        self.fp = fp
73
 
        self.readBytes = 0
74
 
        self.length = length
75
 
        self.finished = 0
76
 
        
77
 
    def tell(self):
78
 
        return self.readBytes
79
 
    
80
 
    def read(self, n=None):
81
 
        if n is None:
82
 
            n = self.length - self.readBytes
83
 
        if n == 0 or self.finished:
84
 
            return ''
85
 
        
86
 
        data = self.fp.read(min(n, self.length - self.readBytes))
87
 
        self.readBytes += len(data)
88
 
        if self.readBytes == self.length or len(data) <  n:
89
 
            self.finished = 1
90
 
        return data
91
 
 
92
 
    def close(self):
93
 
        self.finished = 1
94
 
        del self.fp
95
 
 
96
 
 
97
 
class DeflatedZipFileEntry:
98
 
    """File-like object used to read a deflated entry in a ZipFile"""
99
 
    
100
 
    def __init__(self, fp, length):
101
 
        self.fp = fp
102
 
        self.returnedBytes = 0
103
 
        self.readBytes = 0
104
 
        self.decomp = zlib.decompressobj(-15)
105
 
        self.buffer = ""
106
 
        self.length = length
107
 
        self.finished = 0
108
 
        
109
 
    def tell(self):
110
 
        return self.returnedBytes
111
 
    
112
 
    def read(self, n=None):
113
 
        if self.finished:
114
 
            return ""
115
 
        if n is None:
116
 
            result = [self.buffer,]
117
 
            result.append(self.decomp.decompress(self.fp.read(self.length - self.readBytes)))
118
 
            result.append(self.decomp.decompress("Z"))
119
 
            result.append(self.decomp.flush())
120
 
            self.buffer = ""
121
 
            self.finished = 1
122
 
            result = "".join(result)
123
 
            self.returnedBytes += len(result)
124
 
            return result
125
 
        else:
126
 
            while len(self.buffer) < n:
127
 
                data = self.fp.read(min(n, 1024, self.length - self.readBytes))
128
 
                self.readBytes += len(data)
129
 
                if not data:
130
 
                    result = self.buffer + self.decomp.decompress("Z") + self.decomp.flush()
131
 
                    self.finished = 1
132
 
                    self.buffer = ""
133
 
                    self.returnedBytes += len(result)
134
 
                    return result
135
 
                else:
136
 
                    self.buffer += self.decomp.decompress(data)
137
 
            result = self.buffer[:n]
138
 
            self.buffer = self.buffer[n:]
139
 
            self.returnedBytes += len(result)
140
 
            return result
141
 
    
142
 
    def close(self):
143
 
        self.finished = 1
144
 
        del self.fp
145
 
 
146
 
 
147
 
def unzip(filename, directory=".", overwrite=0):
148
 
    """Unzip the file
149
 
    @param filename: the name of the zip file
150
 
    @param directory: the directory into which the files will be
151
 
    extracted
152
 
    @param overwrite: if on, overwrite files when they exist.  You can
153
 
    still get an error if you try to create a directory over a file
154
 
    with the same name or vice-versa.
155
 
    """
156
 
    for i in unzipIter(filename, directory, overwrite):
157
 
        pass
158
 
 
159
 
DIR_BIT=16
160
 
def unzipIter(filename, directory='.', overwrite=0):
161
 
    """Return a generator for the zipfile.  This implementation will
162
 
    yield after every file.
163
 
 
164
 
    The value it yields is the number of files left to unzip.
165
 
    """
166
 
    zf=zipfile.ZipFile(filename, 'r')
167
 
    names=zf.namelist()
168
 
    if not os.path.exists(directory): os.makedirs(directory)
169
 
    remaining=countZipFileEntries(filename)
170
 
    for entry in names:
171
 
        remaining=remaining - 1
172
 
        isdir=zf.getinfo(entry).external_attr & DIR_BIT
173
 
        f=os.path.join(directory, entry)
174
 
        if isdir:
175
 
            # overwrite flag only applies to files
176
 
            if not os.path.exists(f): os.makedirs(f)
177
 
        else:
178
 
            # create the directory the file will be in first,
179
 
            # since we can't guarantee it exists
180
 
            fdir=os.path.split(f)[0]
181
 
            if not os.path.exists(fdir):
182
 
                os.makedirs(f)
183
 
            if overwrite or not os.path.exists(f):
184
 
                outfile=file(f, 'wb')
185
 
                outfile.write(zf.read(entry))
186
 
                outfile.close()
187
 
        yield remaining
188
 
 
189
 
def countZipFileChunks(filename, chunksize):
190
 
    """Predict the number of chunks that will be extracted from the
191
 
    entire zipfile, given chunksize blocks.
192
 
    """
193
 
    totalchunks=0
194
 
    zf=ChunkingZipFile(filename)
195
 
    for info in zf.infolist():
196
 
        totalchunks=totalchunks+countFileChunks(info, chunksize)
197
 
    return totalchunks
198
 
 
199
 
def countFileChunks(zipinfo, chunksize):
200
 
    size=zipinfo.file_size
201
 
    count=size/chunksize
202
 
    if size%chunksize > 0:
203
 
        count=count+1
204
 
    # each file counts as at least one chunk
205
 
    return count or 1
206
 
    
207
 
def countZipFileEntries(filename):
208
 
    zf=zipfile.ZipFile(filename)
209
 
    return len(zf.namelist())
210
 
 
211
 
def unzipIterChunky(filename, directory='.', overwrite=0,
212
 
                    chunksize=4096):
213
 
    """Return a generator for the zipfile.  This implementation will
214
 
    yield after every chunksize uncompressed bytes, or at the end of a
215
 
    file, whichever comes first.
216
 
 
217
 
    The value it yields is the number of chunks left to unzip.
218
 
    """
219
 
    czf=ChunkingZipFile(filename, 'r')
220
 
    if not os.path.exists(directory): os.makedirs(directory)
221
 
    remaining=countZipFileChunks(filename, chunksize)
222
 
    names=czf.namelist()
223
 
    infos=czf.infolist()
224
 
    
225
 
    for entry, info in zip(names, infos):
226
 
        isdir=info.external_attr & DIR_BIT
227
 
        f=os.path.join(directory, entry)
228
 
        if isdir:
229
 
            # overwrite flag only applies to files
230
 
            if not os.path.exists(f): os.makedirs(f)
231
 
            remaining=remaining-1
232
 
            assert remaining>=0
233
 
            yield remaining
234
 
        else:
235
 
            # create the directory the file will be in first,
236
 
            # since we can't guarantee it exists
237
 
            fdir=os.path.split(f)[0]
238
 
            if not os.path.exists(fdir):
239
 
                os.makedirs(f)
240
 
            if overwrite or not os.path.exists(f):
241
 
                outfile=file(f, 'wb')
242
 
                fp=czf.readfile(entry)
243
 
                if info.file_size==0:
244
 
                    remaining=remaining-1
245
 
                    assert remaining>=0
246
 
                    yield remaining
247
 
                fread=fp.read
248
 
                ftell=fp.tell
249
 
                owrite=outfile.write
250
 
                size=info.file_size
251
 
                while ftell() < size:
252
 
                    hunk=fread(chunksize)
253
 
                    owrite(hunk)
254
 
                    remaining=remaining-1
255
 
                    assert remaining>=0
256
 
                    yield remaining
257
 
                outfile.close()
258
 
            else:
259
 
                remaining=remaining-countFileChunks(info, chunksize)
260
 
                assert remaining>=0
261
 
                yield remaining