added response decompression tool

--HG--
extra : convert_revision : svn%3Ab85faa78-f9eb-468e-a121-7cced6da292c%40213
This commit is contained in:
elpolilla 2008-09-09 12:56:01 +00:00
parent 3a2981801c
commit 7df63477bf
2 changed files with 98 additions and 0 deletions

View File

@ -0,0 +1,45 @@
import os
from unittest import TestCase, main
from scrapy.http import Response, ResponseBody
from scrapy.utils.decompressor import Decompressor
class ScrapyDecompressTest(TestCase):
uncompressed_body = ''
test_responses = {}
decompressor = Decompressor()
def setUp(self):
self.datadir = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'sample_data', 'compressed')
formats = ['tar', 'xml.bz2', 'xml.gz', 'zip']
uncompressed_fd = open(os.path.join(self.datadir, 'feed-sample1.xml'), 'r')
self.uncompressed_body = uncompressed_fd.read()
uncompressed_fd.close()
for format in formats:
fd = open(os.path.join(self.datadir, 'feed-sample1.' + format), 'r')
body = ResponseBody(fd.read())
fd.close()
self.test_responses[format] = Response('foo.com', 'http://foo.com/bar', body=body)
def test_tar(self):
ret = self.decompressor.extract(self.test_responses['tar'])
if ret:
self.assertEqual(ret.body.to_string(), self.uncompressed_body)
def test_zip(self):
ret = self.decompressor.extract(self.test_responses['zip'])
if ret:
self.assertEqual(ret.body.to_string(), self.uncompressed_body)
def test_gz(self):
ret = self.decompressor.extract(self.test_responses['xml.gz'])
if ret:
self.assertEqual(ret.body.to_string(), self.uncompressed_body)
def test_bz2(self):
ret = self.decompressor.extract(self.test_responses['xml.bz2'])
if ret:
self.assertEqual(ret.body.to_string(), self.uncompressed_body)
if __name__ == '__main__':
main()

View File

@ -0,0 +1,53 @@
"""
Utility for autodetecting and decompressing responses
"""
import zipfile
import tarfile
import gzip
import bz2
from tempfile import NamedTemporaryFile
from scrapy.http import ResponseBody
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
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