home *** CD-ROM | disk | FTP | other *** search
/ OS/2 Shareware BBS: 10 Tools / 10-Tools.zip / pyos2bin.zip / Lib / gzip.py < prev    next >
Text File  |  1997-12-30  |  7KB  |  302 lines
  1. import time
  2. import string
  3. import zlib
  4. import StringIO
  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.     t = divmod(value, 256)
  19.     b1 = chr(t[1])
  20.  
  21.     t = divmod(t[0], 256)
  22.     b2 = chr(t[1])
  23.  
  24.     t = divmod(t[0], 256)
  25.     b3 = chr(t[1])
  26.  
  27.     t = divmod(t[0], 256)
  28.     b4 = chr(t[1])
  29.  
  30.     buf = b1 + b2 + b3 + b4
  31.     output.write(buf)
  32.    
  33.     
  34. def read32(input):
  35.     buf = input.read(4)
  36.     v = ord(buf[0])
  37.     v = v + (ord(buf[1]) << 8)
  38.     v = v + (ord(buf[2]) << 16)
  39.     v = v + (ord(buf[3]) << 24)
  40.     return v
  41.  
  42. def open(filename, mode="r", compresslevel=9):
  43.     return GzipFile(filename, mode, compresslevel)
  44.  
  45. class GzipFile:
  46.  
  47.     myfileobj = None
  48.  
  49.     def __init__(self, filename=None, mode=None, 
  50.          compresslevel=9, fileobj=None):
  51.     if fileobj is None:
  52.         fileobj = self.myfileobj = __builtin__.open(filename, mode or 'r')
  53.         if filename is None:
  54.         if hasattr(fileobj, 'name'): filename = fileobj.name
  55.         else: filename = ''
  56.         if mode is None:
  57.         if hasattr(fileobj, 'mode'): mode = fileobj.mode
  58.         else: mode = 'r'
  59.  
  60.     if mode[0:1] == 'r':
  61.         self.mode = READ
  62.         self._init_read()
  63.         self.filename = filename
  64.         self.decompress = zlib.decompressobj(-zlib.MAX_WBITS)
  65.  
  66.     elif mode[0:1] == 'w':
  67.         self.mode = WRITE
  68.         self._init_write(filename)
  69.         self.compress = zlib.compressobj(compresslevel,
  70.                          zlib.DEFLATED, 
  71.                          -zlib.MAX_WBITS,
  72.                          zlib.DEF_MEM_LEVEL,
  73.                          0)
  74.     else:
  75.         raise ValueError, "Mode " + mode + " not supported"
  76.  
  77.     self.fileobj = fileobj
  78.  
  79.     if self.mode == WRITE:
  80.         self._write_gzip_header()
  81.     elif self.mode == READ:
  82.         self._read_gzip_header()
  83.         
  84.  
  85.     def __repr__(self):
  86.     s = repr(self.fileobj)
  87.     return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>'
  88.  
  89.     def _init_write(self, filename):
  90.     if filename[-3:] != '.gz':
  91.         filename = filename + '.gz'
  92.     self.filename = filename
  93.     self.crc = zlib.crc32("")
  94.     self.size = 0
  95.     self.writebuf = []
  96.     self.bufsize = 0
  97.  
  98.     def _write_gzip_header(self):
  99.     self.fileobj.write('\037\213')             # magic header
  100.     self.fileobj.write('\010')                 # compression method
  101.     fname = self.filename[:-3]
  102.     flags = 0
  103.     if fname:
  104.         flags = FNAME
  105.     self.fileobj.write(chr(flags))
  106.     write32(self.fileobj, int(time.time()))
  107.     self.fileobj.write('\002')
  108.     self.fileobj.write('\377')
  109.     if fname:
  110.         self.fileobj.write(fname + '\000')
  111.  
  112.     def _init_read(self):
  113.     self.crc = zlib.crc32("")
  114.     self.size = 0
  115.     self.extrabuf = ""
  116.     self.extrasize = 0
  117.  
  118.     def _read_gzip_header(self):
  119.     magic = self.fileobj.read(2)
  120.     if magic != '\037\213':
  121.         raise RuntimeError, 'Not a gzipped file'
  122.     method = ord( self.fileobj.read(1) )
  123.     if method != 8:
  124.         raise RuntimeError, 'Unknown compression method'
  125.     flag = ord( self.fileobj.read(1) )
  126.     # modtime = self.fileobj.read(4)
  127.     # extraflag = self.fileobj.read(1)
  128.     # os = self.fileobj.read(1)
  129.     self.fileobj.read(6)
  130.  
  131.     if flag & FEXTRA:
  132.         # Read & discard the extra field, if present
  133.         xlen=ord(self.fileobj.read(1))        
  134.         xlen=xlen+256*ord(self.fileobj.read(1))
  135.         self.fileobj.read(xlen)
  136.     if flag & FNAME:
  137.         # Read and discard a null-terminated string containing the filename
  138.         while (1):
  139.         s=self.fileobj.read(1)
  140.         if not s or s=='\000': break
  141.     if flag & FCOMMENT:
  142.         # Read and discard a null-terminated string containing a comment
  143.         while (1):
  144.         s=self.fileobj.read(1)
  145.         if not s or s=='\000': break
  146.     if flag & FHCRC:
  147.         self.fileobj.read(2)     # Read & discard the 16-bit header CRC
  148.  
  149.  
  150.     def write(self,data):
  151.     if self.fileobj is None:
  152.         raise ValueError, "write() on closed GzipFile object"
  153.     if len(data) > 0:
  154.         self.size = self.size + len(data)
  155.         self.crc = zlib.crc32(data, self.crc)
  156.         self.fileobj.write( self.compress.compress(data) )
  157.  
  158.     def writelines(self,lines):
  159.     self.write(string.join(lines))
  160.  
  161.     def read(self,size=None):
  162.     if self.extrasize <= 0 and self.fileobj is None:
  163.         return ''
  164.     
  165.     if not size:
  166.         # get the whole thing
  167.         try:
  168.         while 1:
  169.             self._read()
  170.         except EOFError:
  171.         size = self.extrasize
  172.     else:
  173.         # just get some more of it
  174.         try:
  175.         while size > self.extrasize:
  176.             self._read()
  177.         except EOFError:
  178.         pass
  179.     
  180.     chunk = self.extrabuf[:size]
  181.     self.extrabuf = self.extrabuf[size:]
  182.     self.extrasize = self.extrasize - size
  183.  
  184.     return chunk
  185.  
  186.     def _read(self):
  187.     buf = self.fileobj.read(1024)
  188.     if buf == "":
  189.         uncompress = self.decompress.flush()
  190.         if uncompress == "":
  191.         self._read_eof()
  192.         self.fileobj = None
  193.         raise EOFError, 'Reached EOF'
  194.     else:
  195.         uncompress = self.decompress.decompress(buf)
  196.     self.crc = zlib.crc32(uncompress, self.crc)
  197.     self.extrabuf = self.extrabuf + uncompress
  198.     self.extrasize = self.extrasize + len(uncompress)
  199.     self.size = self.size + len(uncompress)
  200.  
  201.     def _read_eof(self):
  202.     # Andrew writes:
  203.     ## We've read to the end of the file, so we have to rewind in order
  204.     ## to reread the 8 bytes containing the CRC and the file size.  The
  205.     ## decompressor is smart and knows when to stop, so feeding it
  206.     ## extra data is harmless.  
  207.     self.fileobj.seek(-8, 2)
  208.     crc32 = read32(self.fileobj)
  209.     isize = read32(self.fileobj)
  210.     if crc32 != self.crc:
  211.         self.error = "CRC check failed"
  212.     elif isize != self.size:
  213.         self.error = "Incorrect length of data produced"
  214.  
  215.     def close(self):
  216.     if self.mode == WRITE:
  217.         self.fileobj.write(self.compress.flush())
  218.         write32(self.fileobj, self.crc)
  219.         write32(self.fileobj, self.size)
  220.         self.fileobj = None
  221.     elif self.mode == READ:
  222.         self.fileobj = None
  223.     if self.myfileobj:
  224.         self.myfileobj.close()
  225.         self.myfileobj = None
  226.  
  227.     def flush(self):
  228.     self.fileobj.flush()
  229.  
  230.     def seek(self):
  231.     raise IOError, 'Random access not allowed in gzip files'
  232.  
  233.     def tell(self):
  234.     raise IOError, 'I won\'t tell() you for gzip files'
  235.  
  236.     def isatty(self):
  237.     return 0
  238.  
  239.     def readline(self):
  240.     # XXX This function isn't implemented in a very efficient way
  241.     line=""
  242.     while 1:
  243.         c = self.read(1)
  244.         line = line + c
  245.         if c=='\n' or c=="": break
  246.     return line
  247.  
  248.     def readlines(self):
  249.     L=[]
  250.     line = self.readline()
  251.     while line!="":
  252.         L.append(line)
  253.         line = self.readline()
  254.     return L
  255.  
  256.     def writelines(self, L):
  257.     for line in L:
  258.         self.write(line)
  259.  
  260.  
  261. def _test():
  262.     # Act like gzip; with -d, act like gunzip.
  263.     # The input file is not deleted, however, nor are any other gzip
  264.     # options or features supported.
  265.     import sys
  266.     args = sys.argv[1:]
  267.     decompress = args and args[0] == "-d"
  268.     if decompress:
  269.     args = args[1:]
  270.     if not args:
  271.     args = ["-"]
  272.     for arg in args:
  273.     if decompress:
  274.         if arg == "-":
  275.         f = GzipFile(filename="", mode="rb", fileobj=sys.stdin)
  276.         g = sys.stdout
  277.         else:
  278.         if arg[-3:] != ".gz":
  279.             print "filename doesn't end in .gz:", `arg`
  280.             continue
  281.         f = open(arg, "rb")
  282.         g = __builtin__.open(arg[:-3], "wb")
  283.     else:
  284.         if arg == "-":
  285.         f = sys.stdin
  286.         g = GzipFile(filename="", mode="wb", fileobj=sys.stdout)
  287.         else:
  288.         f = __builtin__.open(arg, "rb")
  289.         g = open(arg + ".gz", "wb")
  290.     while 1:
  291.         chunk = f.read(1024)
  292.         if not chunk:
  293.         break
  294.         g.write(chunk)
  295.     if g is not sys.stdout:
  296.         g.close()
  297.     if f is not sys.stdin:
  298.         f.close()
  299.  
  300. if __name__ == '__main__':
  301.     _test()
  302.