view zipfiles/zipfile314.diff @ 217:ce1285728952

Fixed the signature of @abstractmethod TestCase.test
author Oleg Oshmyan <chortos@inbox.lv>
date Sun, 21 Aug 2011 01:24:29 +0300
parents b993d9257400
children
line wrap: on
line source

--- /usr/local/lib/python3.1/zipfile.py	2011-06-16 01:02:30.000000000 +0100
+++ zipfile31.py	2011-06-16 01:16:53.000000000 +0100
@@ -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 @@ class LargeZipFile(Exception):
 # 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 @@     def __init__(self, fileobj, zipinfo,
         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 @@     def read(self, size = None):
             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 @@     def read(self, size = None):
                     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 @@ class ZipFile:
     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 @@     def __init__(self, file, mode="r", c
             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")
 
@@ -1053,7 +1071,10 @@     def _writecheck(self, zinfo):
         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:
@@ -1114,6 +1135,8 @@     def write(self, filename, arcname=No
             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:
@@ -1174,6 +1197,10 @@     def writestr(self, zinfo_or_arcname,
                  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