diff --git a/scrapy/trunk/scrapy/utils/decompressor.py b/scrapy/trunk/scrapy/utils/decompressor.py index 60334f3b7..3b951038c 100644 --- a/scrapy/trunk/scrapy/utils/decompressor.py +++ b/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