mirror of https://github.com/scrapy/scrapy.git
added response decompression tool
--HG-- extra : convert_revision : svn%3Ab85faa78-f9eb-468e-a121-7cced6da292c%40213
This commit is contained in:
parent
3a2981801c
commit
7df63477bf
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
||||
Loading…
Reference in New Issue