home *** CD-ROM | disk | FTP | other *** search
/ OS/2 Shareware BBS: 10 Tools / 10-Tools.zip / pytho152.zip / emx / lib / python1.5 / gzip.py < prev    next >
Text File  |  2000-08-10  |  11KB  |  334 lines
  1. import time
  2. import string
  3. import zlib
  4. import struct
  5. import __builtin__
  6.  
  7. # implements a python function that reads and writes a gzipped file
  8. # the user of the file doesn't have to worry about the compression,
  9. # but random access is not allowed
  10.  
  11. # based on Andrew Kuchling's minigzip.py distributed with the zlib module
  12.  
  13. FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16
  14.  
  15. READ, WRITE = 1, 2
  16.  
  17. def write32(output, value):
  18.     output.write(struct.pack("<l", value))
  19.     
  20. def write32u(output, value):
  21.     output.write(struct.pack("<L", value))
  22.  
  23. def read32(input):
  24.     return struct.unpack("<l", input.read(4))[0]
  25.  
  26. def open(filename, mode="rb", compresslevel=9):
  27.     return GzipFile(filename, mode, compresslevel)
  28.  
  29. class GzipFile:
  30.  
  31.     myfileobj = None
  32.  
  33.     def __init__(self, filename=None, mode=None, 
  34.                  compresslevel=9, fileobj=None):
  35.         if fileobj is None:
  36.             fileobj = self.myfileobj = __builtin__.open(filename, mode or 'rb')
  37.         if filename is None:
  38.             if hasattr(fileobj, 'name'): filename = fileobj.name
  39.             else: filename = ''
  40.         if mode is None:
  41.             if hasattr(fileobj, 'mode'): mode = fileobj.mode
  42.             else: mode = 'rb'
  43.  
  44.         if mode[0:1] == 'r':
  45.             self.mode = READ
  46.          # Set flag indicating start of a new member
  47.             self._new_member = 1 
  48.             self.extrabuf = ""
  49.             self.extrasize = 0
  50.             self.filename = filename
  51.  
  52.         elif mode[0:1] == 'w' or mode[0:1] == 'a':
  53.             self.mode = WRITE
  54.             self._init_write(filename)
  55.             self.compress = zlib.compressobj(compresslevel,
  56.                                              zlib.DEFLATED, 
  57.                                              -zlib.MAX_WBITS,
  58.                                              zlib.DEF_MEM_LEVEL,
  59.                                              0)
  60.         else:
  61.             raise ValueError, "Mode " + mode + " not supported"
  62.  
  63.         self.fileobj = fileobj
  64.  
  65.         if self.mode == WRITE:
  66.             self._write_gzip_header()
  67.  
  68.     def __repr__(self):
  69.         s = repr(self.fileobj)
  70.         return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>'
  71.  
  72.     def _init_write(self, filename):
  73.         if filename[-3:] != '.gz':
  74.             filename = filename + '.gz'
  75.         self.filename = filename
  76.         self.crc = zlib.crc32("")
  77.         self.size = 0
  78.         self.writebuf = []
  79.         self.bufsize = 0
  80.  
  81.     def _write_gzip_header(self):
  82.         self.fileobj.write('\037\213')             # magic header
  83.         self.fileobj.write('\010')                 # compression method
  84.         fname = self.filename[:-3]
  85.         flags = 0
  86.         if fname:
  87.             flags = FNAME
  88.         self.fileobj.write(chr(flags))
  89.         write32u(self.fileobj, long(time.time()))
  90.         self.fileobj.write('\002')
  91.         self.fileobj.write('\377')
  92.         if fname:
  93.             self.fileobj.write(fname + '\000')
  94.  
  95.     def _init_read(self):
  96.         self.crc = zlib.crc32("")
  97.         self.size = 0
  98.  
  99.     def _read_gzip_header(self):
  100.         magic = self.fileobj.read(2)
  101.         if magic != '\037\213':
  102.             raise IOError, 'Not a gzipped file'
  103.         method = ord( self.fileobj.read(1) )
  104.         if method != 8:
  105.             raise IOError, 'Unknown compression method'
  106.         flag = ord( self.fileobj.read(1) )
  107.         # modtime = self.fileobj.read(4)
  108.         # extraflag = self.fileobj.read(1)
  109.         # os = self.fileobj.read(1)
  110.         self.fileobj.read(6)
  111.  
  112.         if flag & FEXTRA:
  113.             # Read & discard the extra field, if present
  114.             xlen=ord(self.fileobj.read(1))              
  115.             xlen=xlen+256*ord(self.fileobj.read(1))
  116.             self.fileobj.read(xlen)
  117.         if flag & FNAME:
  118.             # Read and discard a null-terminated string containing the filename
  119.             while (1):
  120.                 s=self.fileobj.read(1)
  121.                 if not s or s=='\000': break
  122.         if flag & FCOMMENT:
  123.             # Read and discard a null-terminated string containing a comment
  124.             while (1):
  125.                 s=self.fileobj.read(1)
  126.                 if not s or s=='\000': break
  127.         if flag & FHCRC:
  128.             self.fileobj.read(2)     # Read & discard the 16-bit header CRC
  129.  
  130.  
  131.     def write(self,data):
  132.         if self.fileobj is None:
  133.             raise ValueError, "write() on closed GzipFile object"
  134.         if len(data) > 0:
  135.             self.size = self.size + len(data)
  136.             self.crc = zlib.crc32(data, self.crc)
  137.             self.fileobj.write( self.compress.compress(data) )
  138.  
  139.     def writelines(self,lines):
  140.         self.write(string.join(lines))
  141.  
  142.     def read(self, size=None):
  143.         if self.extrasize <= 0 and self.fileobj is None:
  144.             return ''
  145.  
  146.         readsize = 1024
  147.         if not size:        # get the whole thing
  148.             try:
  149.                 while 1:
  150.                     self._read(readsize)
  151.                     readsize = readsize * 2
  152.             except EOFError:
  153.                 size = self.extrasize
  154.         else:               # just get some more of it
  155.             try:
  156.                 while size > self.extrasize:
  157.                     self._read(readsize)
  158.                     readsize = readsize * 2
  159.             except EOFError:
  160.                 if size > self.extrasize:
  161.                     size = self.extrasize
  162.         
  163.         chunk = self.extrabuf[:size]
  164.         self.extrabuf = self.extrabuf[size:]
  165.         self.extrasize = self.extrasize - size
  166.  
  167.         return chunk
  168.  
  169.     def _unread(self, buf):
  170.         self.extrabuf = buf + self.extrabuf
  171.         self.extrasize = len(buf) + self.extrasize
  172.  
  173.     def _read(self, size=1024):
  174.         if self.fileobj is None: raise EOFError, "Reached EOF"
  175.      
  176.         if self._new_member:
  177.             # If the _new_member flag is set, we have to 
  178.             # 
  179.             # First, check if we're at the end of the file;
  180.             # if so, it's time to stop; no more members to read.
  181.             pos = self.fileobj.tell()   # Save current position
  182.             self.fileobj.seek(0, 2)     # Seek to end of file
  183.             if pos == self.fileobj.tell():
  184.                 self.fileobj = None
  185.                 return EOFError, "Reached EOF"
  186.             else: 
  187.                 self.fileobj.seek( pos ) # Return to original position
  188.   
  189.             self._init_read()       
  190.             self._read_gzip_header()
  191.             self.decompress = zlib.decompressobj(-zlib.MAX_WBITS)
  192.             self._new_member = 0
  193.  
  194.         # Read a chunk of data from the file
  195.         buf = self.fileobj.read(size)
  196.  
  197.         # If the EOF has been reached, flush the decompression object
  198.         # and mark this object as finished.
  199.        
  200.         if buf == "":
  201.             uncompress = self.decompress.flush()
  202.             self._read_eof()
  203.             self.fileobj = None
  204.             self._add_read_data( uncompress )
  205.             raise EOFError, 'Reached EOF'
  206.   
  207.         uncompress = self.decompress.decompress(buf)
  208.         self._add_read_data( uncompress )
  209.  
  210.         if self.decompress.unused_data != "":
  211.             # Ending case: we've come to the end of a member in the file,
  212.             # so seek back to the start of the unused data, finish up
  213.             # this member, and read a new gzip header.
  214.             # (The number of bytes to seek back is the length of the unused
  215.             # data, minus 8 because _read_eof() will rewind a further 8 bytes)
  216.             self.fileobj.seek( -len(self.decompress.unused_data)+8, 1)
  217.  
  218.             # Check the CRC and file size, and set the flag so we read
  219.             # a new member on the next call 
  220.             self._read_eof()
  221.             self._new_member = 1        
  222.         
  223.     def _add_read_data(self, data):            
  224.         self.crc = zlib.crc32(data, self.crc)
  225.         self.extrabuf = self.extrabuf + data
  226.         self.extrasize = self.extrasize + len(data)
  227.         self.size = self.size + len(data)
  228.  
  229.     def _read_eof(self):
  230.         # We've read to the end of the file, so we have to rewind in order
  231.         # to reread the 8 bytes containing the CRC and the file size.  
  232.         # We check the that the computed CRC and size of the
  233.         # uncompressed data matches the stored values.
  234.         self.fileobj.seek(-8, 1)
  235.         crc32 = read32(self.fileobj)
  236.         isize = read32(self.fileobj)
  237.         if crc32%0x100000000L != self.crc%0x100000000L:
  238.             raise ValueError, "CRC check failed"
  239.         elif isize != self.size:
  240.             raise ValueError, "Incorrect length of data produced"
  241.           
  242.     def close(self):
  243.         if self.mode == WRITE:
  244.             self.fileobj.write(self.compress.flush())
  245.             write32(self.fileobj, self.crc)
  246.             write32(self.fileobj, self.size)
  247.             self.fileobj = None
  248.         elif self.mode == READ:
  249.             self.fileobj = None
  250.         if self.myfileobj:
  251.             self.myfileobj.close()
  252.             self.myfileobj = None
  253.  
  254.     def flush(self):
  255.         self.fileobj.flush()
  256.  
  257.     def seek(self):
  258.         raise IOError, 'Random access not allowed in gzip files'
  259.  
  260.     def tell(self):
  261.         raise IOError, 'I won\'t tell() you for gzip files'
  262.  
  263.     def isatty(self):
  264.         return 0
  265.  
  266.     def readline(self):
  267.         bufs = []
  268.         readsize = 100
  269.         while 1:
  270.             c = self.read(readsize)
  271.             i = string.find(c, '\n')
  272.             if i >= 0 or c == '':
  273.                 bufs.append(c[:i+1])
  274.                 self._unread(c[i+1:])
  275.                 return string.join(bufs, '')
  276.             bufs.append(c)
  277.             readsize = readsize * 2
  278.  
  279.     def readlines(self):
  280.         buf = self.read()
  281.         lines = string.split(buf, '\n')
  282.         for i in range(len(lines)-1):
  283.             lines[i] = lines[i] + '\n'
  284.         if lines and not lines[-1]:
  285.             del lines[-1]
  286.         return lines
  287.  
  288.     def writelines(self, L):
  289.         for line in L:
  290.             self.write(line)
  291.  
  292.  
  293. def _test():
  294.     # Act like gzip; with -d, act like gunzip.
  295.     # The input file is not deleted, however, nor are any other gzip
  296.     # options or features supported.
  297.     import sys
  298.     args = sys.argv[1:]
  299.     decompress = args and args[0] == "-d"
  300.     if decompress:
  301.         args = args[1:]
  302.     if not args:
  303.         args = ["-"]
  304.     for arg in args:
  305.         if decompress:
  306.             if arg == "-":
  307.                 f = GzipFile(filename="", mode="rb", fileobj=sys.stdin)
  308.                 g = sys.stdout
  309.             else:
  310.                 if arg[-3:] != ".gz":
  311.                     print "filename doesn't end in .gz:", `arg`
  312.                     continue
  313.                 f = open(arg, "rb")
  314.                 g = __builtin__.open(arg[:-3], "wb")
  315.         else:
  316.             if arg == "-":
  317.                 f = sys.stdin
  318.                 g = GzipFile(filename="", mode="wb", fileobj=sys.stdout)
  319.             else:
  320.                 f = __builtin__.open(arg, "rb")
  321.                 g = open(arg + ".gz", "wb")
  322.         while 1:
  323.             chunk = f.read(1024)
  324.             if not chunk:
  325.                 break
  326.             g.write(chunk)
  327.         if g is not sys.stdout:
  328.             g.close()
  329.         if f is not sys.stdin:
  330.             f.close()
  331.  
  332. if __name__ == '__main__':
  333.     _test()
  334.