decompressor tool improved

--HG--
extra : convert_revision : svn%3Ab85faa78-f9eb-468e-a121-7cced6da292c%40216

scrapy/trunk/scrapy/utils/decompressor.py

@@ -6,6 +6,7 @@ import zipfile
 import tarfile
 import gzip
 import bz2
+from cStringIO import StringIO
 from tempfile import NamedTemporaryFile
 from scrapy.http import ResponseBody
 
@@ -13,41 +14,54 @@ class Decompressor(object):
     class ArchiveIsEmpty(Exception):
         pass
     
-    def extract(self, response):
-        temp = NamedTemporaryFile()
-        temp.file.write(response.body.to_string())
-        temp.file.seek(0)
-        
-        if tarfile.is_tarfile(temp.name):
-            tar = tarfile.open(temp.name)
-            if tar.members:
-                return response.replace(body=ResponseBody(tar.extractfile(tar.members[0]).read()))
-            else:
-                raise self.ArchiveIsEmpty
-           
-        elif zipfile.is_zipfile(temp.name):
-            zipf = zipfile.ZipFile(temp.name, 'r')
-            namelist = zipf.namelist()
-            if namelist:
-                return response.replace(body=ResponseBody(zipf.read(namelist[0])))
-            else:
-                raise self.ArchiveIsEmpty
-           
+    def __init__(self):
+        self.decompressors = [self.is_tar, self.is_zip,
+                              self.is_gzip, self.is_bzip2]        
+    def is_tar(self, response):
+        try:
+            tar_file = tarfile.open(fileobj=self.archive)
+        except tarfile.ReadError:
+            return False
+        if tar_file.members:
+            return response.replace(body=ResponseBody(tar_file.extractfile(tar_file.members[0]).read()))
         else:
-            # It's neither a tar or a zip, so we try to decompress using Gzip now
-            try:
-                gzip_file = gzip.GzipFile(temp.name)
-                return response.replace(body=ResponseBody(gzip_file.read()))
-            except IOError:
-                pass
-      
-            # Finally, we try with Bzip2
-            try:
-                bzip_file = bz2.BZ2File(temp.name)
-                return response.replace(body=ResponseBody(bzip_file.read()))
-            except IOError:
-                pass
-        
-            # We couldn't decompress the file, so we return the same response
-            return response
+            raise self.ArchiveIsEmpty
     
+    def is_zip(self, response):
+        try:
+            zip_file = zipfile.ZipFile(self.archive)
+        except zipfile.BadZipfile:
+            return False
+        namelist = zip_file.namelist()
+        if namelist:
+            return response.replace(body=ResponseBody(zip_file.read(namelist[0])))
+        else:
+            raise self.ArchiveIsEmpty
+            
+    def is_gzip(self, response):
+        try:
+            gzip_file = gzip.GzipFile(fileobj=self.archive)
+            decompressed_body = gzip_file.read()
+        except IOError:
+            return False
+        return response.replace(body=decompressed_body)
+            
+    def is_bzip2(self, response):
+        try:
+            decompressed_body = bz2.decompress(self.body)
+        except IOError:
+            return False
+        return response.replace(body=ResponseBody(decompressed_body))
+            
+    def extract(self, response):
+        self.body = response.body.to_string()
+        self.archive = StringIO()
+        self.archive.write(self.body)
+        self.archive.seek(0)
+
+        for decompressor in self.decompressors:
+            self.archive.seek(0)
+            ret = decompressor(response)
+            if ret:
+                return ret
+        return response