view zipfile313.diff @ 99:9b41934b2e95

Fixed a major regression introduced in changeset 1fb319ec33af testcases.load_problem now always handles the True sent by Problem.test.
author Oleg Oshmyan <chortos@inbox.lv>
date Wed, 02 Mar 2011 20:03:44 +0000
parents 4ea7133ac25c
children
line wrap: on
line source

--- /usr/local/lib/python3.1/zipfile.py	2010-11-29 00:59:28.000000000 +0000
+++ zipfile313.py	2010-11-29 01:22:19.000000000 +0000
@@ -3,6 +3,7 @@
 
 XXX references to utf-8 need further investigation.
 """
+# Improved by Chortos-2 in 2010 (added bzip2 support)
 import struct, os, time, sys, shutil
 import binascii, io, stat
 
@@ -13,8 +14,13 @@
     zlib = None
     crc32 = binascii.crc32
 
+try:
+    import bz2 # We may need its compression method
+except ImportError:
+    bz2 = None
+
 __all__ = ["BadZipfile", "error", "ZIP_STORED", "ZIP_DEFLATED", "is_zipfile",
-           "ZipInfo", "ZipFile", "PyZipFile", "LargeZipFile" ]
+           "ZipInfo", "ZipFile", "PyZipFile", "LargeZipFile", "ZIP_BZIP2" ]
 
 class BadZipfile(Exception):
     pass
@@ -35,6 +41,7 @@
 # constants for Zip file compression methods
 ZIP_STORED = 0
 ZIP_DEFLATED = 8
+ZIP_BZIP2 = 12
 # Other ZIP compression methods not supported
 
 # Below are some formats and associated data for reading/writing headers using
@@ -477,6 +484,9 @@
         self.compreadsize = 64*1024
         if self.compress_type == ZIP_DEFLATED:
             self.dc = zlib.decompressobj(-15)
+        elif self.compress_type == ZIP_BZIP2:
+            self.dc = bz2.BZ2Decompressor()
+            self.compreadsize = 900000
 
         if hasattr(zipinfo, 'CRC'):
             self._expected_crc = zipinfo.CRC
@@ -604,7 +614,7 @@
             if self.compress_type == ZIP_STORED:
                 lr = len(self.readbuffer)
                 bytesToRead = min(bytesToRead, size - lr)
-            elif self.compress_type == ZIP_DEFLATED:
+            else:
                 if len(self.readbuffer) > size:
                     # the user has requested fewer bytes than we've already
                     # pulled through the decompressor; don't read any more
@@ -639,14 +649,17 @@
                     newdata = bytes(map(self.decrypter, newdata))
 
                 # decompress newly read data if necessary
-                if newdata and self.compress_type == ZIP_DEFLATED:
+                if newdata and self.compress_type != ZIP_STORED:
                     newdata = self.dc.decompress(newdata)
-                    self.rawbuffer = self.dc.unconsumed_tail
+                    self.rawbuffer = self.dc.unconsumed_tail if self.compress_type == ZIP_DEFLATED else ''
                     if self.eof and len(self.rawbuffer) == 0:
                         # we're out of raw bytes (both from the file and
                         # the local buffer); flush just to make sure the
                         # decompressor is done
-                        newdata += self.dc.flush()
+                        try:
+                            newdata += self.dc.flush()
+                        except AttributeError:
+                            pass
                         # prevent decompressor from being used again
                         self.dc = None
 
@@ -674,7 +687,8 @@
     file: Either the path to the file, or a file-like object.
           If it is a path, the file will be opened and closed by ZipFile.
     mode: The mode can be either read "r", write "w" or append "a".
-    compression: ZIP_STORED (no compression) or ZIP_DEFLATED (requires zlib).
+    compression: ZIP_STORED (no compression), ZIP_DEFLATED (requires zlib),
+                 or ZIP_BZIP2 (requires bz2).
     allowZip64: if True ZipFile will create files with ZIP64 extensions when
                 needed, otherwise it will raise an exception when this would
                 be necessary.
@@ -694,6 +708,10 @@
             if not zlib:
                 raise RuntimeError(
                       "Compression requires the (missing) zlib module")
+        elif compression == ZIP_BZIP2:
+            if not bz2:
+                raise RuntimeError(
+                      "Compression requires the (missing) bz2 module")
         else:
             raise RuntimeError("That compression method is not supported")
 
@@ -1041,7 +1059,10 @@
         if zinfo.compress_type == ZIP_DEFLATED and not zlib:
             raise RuntimeError(
                   "Compression requires the (missing) zlib module")
-        if zinfo.compress_type not in (ZIP_STORED, ZIP_DEFLATED):
+        if zinfo.compress_type == ZIP_BZIP2 and not bz2:
+            raise RuntimeError(
+                  "Compression requires the (missing) bz2 module")
+        if zinfo.compress_type not in (ZIP_STORED, ZIP_DEFLATED, ZIP_BZIP2):
             raise RuntimeError("That compression method is not supported")
         if zinfo.file_size > ZIP64_LIMIT:
             if not self._allowZip64:
@@ -1102,6 +1123,8 @@
             if zinfo.compress_type == ZIP_DEFLATED:
                 cmpr = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
                      zlib.DEFLATED, -15)
+            elif zinfo.compress_type == ZIP_BZIP2:
+                cmpr = bz2.BZ2Compressor()
             else:
                 cmpr = None
             while 1:
@@ -1162,6 +1185,10 @@
                  zlib.DEFLATED, -15)
             data = co.compress(data) + co.flush()
             zinfo.compress_size = len(data)    # Compressed size
+        elif zinfo.compress_type == ZIP_BZIP2:
+            co = bz2.BZ2Compressor()
+            data = co.compress(data) + co.flush()
+            zinfo.compress_size = len(data)    # Compressed size
         else:
             zinfo.compress_size = zinfo.file_size
         zinfo.header_offset = self.fp.tell()    # Start of header data