home *** CD-ROM | disk | FTP | other *** search
/ PC Welt 2006 November (DVD) / PCWELT_11_2006.ISO / casper / filesystem.squashfs / usr / lib / python2.4 / site-packages / BitTorrent / zurllib.py < prev    next >
Encoding:
Python Source  |  2006-07-20  |  4.4 KB  |  154 lines
  1. #
  2. # zurllib.py
  3. #
  4. # This is (hopefully) a drop-in for urllib which will request gzip/deflate
  5. # compression and then decompress the output if a compressed response is
  6. # received while maintaining the API.
  7. #
  8. # by Robert Stone 2/22/2003 
  9. #
  10.  
  11. from urllib import *
  12. from urllib2 import *
  13. from gzip import GzipFile
  14. from StringIO import StringIO
  15. from __init__ import version
  16. import pprint
  17.  
  18.  
  19. DEBUG=0
  20.  
  21.  
  22. class HTTPContentEncodingHandler(HTTPHandler):
  23.     """Inherit and add gzip/deflate/etc support to HTTP gets."""
  24.     def http_open(self, req):
  25.         # add the Accept-Encoding header to the request
  26.         # support gzip encoding (identity is assumed)
  27.         req.add_header("Accept-Encoding","gzip")
  28.         req.add_header('User-Agent', 'BitTorrent/' + version)
  29.         if DEBUG: 
  30.             print "Sending:" 
  31.             print req.headers
  32.             print "\n"
  33.         fp = HTTPHandler.http_open(self,req)
  34.         headers = fp.headers
  35.         if DEBUG: 
  36.              pprint.pprint(headers.dict)
  37.         url = fp.url
  38.     resp = addinfourldecompress(fp, headers, url)
  39.     # As of Python 2.4 http_open response also has 'code' and 'msg'
  40.         # members, and HTTPErrorProcessor breaks if they don't exist.
  41.     if 'code' in dir(fp):
  42.         resp.code = fp.code
  43.     if 'msg' in dir(fp):
  44.         resp.msg = fp.msg
  45.     return resp
  46.  
  47. class addinfourldecompress(addinfourl):
  48.     """Do gzip decompression if necessary. Do addinfourl stuff too."""
  49.     def __init__(self, fp, headers, url):
  50.         # we need to do something more sophisticated here to deal with
  51.         # multiple values?  What about other weird crap like q-values?
  52.         # basically this only works for the most simplistic case and will
  53.         # break in some other cases, but for now we only care about making
  54.         # this work with the BT tracker so....
  55.         if headers.has_key('content-encoding') and headers['content-encoding'] == 'gzip':
  56.             if DEBUG:
  57.                 print "Contents of Content-encoding: " + headers['Content-encoding'] + "\n"
  58.             self.gzip = 1
  59.             self.rawfp = fp
  60.             fp = GzipStream(fp)
  61.         else:
  62.             self.gzip = 0
  63.         return addinfourl.__init__(self, fp, headers, url)
  64.  
  65.     def close(self):
  66.         self.fp.close()
  67.         if self.gzip:
  68.             self.rawfp.close()
  69.  
  70.     def iscompressed(self):
  71.         return self.gzip
  72.  
  73. class GzipStream(StringIO):
  74.     """Magically decompress a file object.
  75.  
  76.        This is not the most efficient way to do this but GzipFile() wants
  77.        to seek, etc, which won't work for a stream such as that from a socket.
  78.        So we copy the whole shebang info a StringIO object, decompress that
  79.        then let people access the decompressed output as a StringIO object.
  80.  
  81.        The disadvantage is memory use and the advantage is random access.
  82.  
  83.        Will mess with fixing this later.
  84.     """
  85.  
  86.     def __init__(self,fp):
  87.         self.fp = fp
  88.  
  89.         # this is nasty and needs to be fixed at some point
  90.         # copy everything into a StringIO (compressed)
  91.         compressed = StringIO()
  92.         r = fp.read()
  93.         while r:
  94.             compressed.write(r)
  95.             r = fp.read()
  96.         # now, unzip (gz) the StringIO to a string
  97.         compressed.seek(0,0)
  98.         gz = GzipFile(fileobj = compressed)
  99.         str = ''
  100.         r = gz.read()
  101.         while r:
  102.             str += r
  103.             r = gz.read()
  104.         # close our utility files
  105.         compressed.close()
  106.         gz.close()
  107.         # init our stringio selves with the string 
  108.         StringIO.__init__(self, str)
  109.         del str
  110.  
  111.     def close(self):
  112.         self.fp.close()
  113.         return StringIO.close(self)
  114.  
  115.  
  116. def test():
  117.     """Test this module.
  118.  
  119.        At the moment this is lame.
  120.     """
  121.  
  122.     print "Running unit tests.\n"
  123.  
  124.     def printcomp(fp):
  125.         try:
  126.             if fp.iscompressed():
  127.                 print "GET was compressed.\n"
  128.             else:
  129.                 print "GET was uncompressed.\n"
  130.         except:
  131.             print "no iscompressed function!  this shouldn't happen"
  132.  
  133.     print "Trying to GET a compressed document...\n"
  134.     fp = urlopen('http://a.scarywater.net/hng/index.shtml')
  135.     print fp.read()
  136.     printcomp(fp)
  137.     fp.close()
  138.  
  139.     print "Trying to GET an unknown document...\n"
  140.     fp = urlopen('http://www.otaku.org/')
  141.     print fp.read()
  142.     printcomp(fp)
  143.     fp.close()
  144.  
  145.  
  146. #
  147. # Install the HTTPContentEncodingHandler that we've defined above.
  148. #
  149. install_opener(build_opener(HTTPContentEncodingHandler))
  150.  
  151. if __name__ == '__main__':
  152.     test()
  153.  
  154.