From 7df63477bff7ccae129fef87d4bf150b6222ba0e Mon Sep 17 00:00:00 2001
From: elpolilla <none@none>
Date: Tue, 9 Sep 2008 12:56:01 +0000
Subject: [PATCH] added response decompression tool

--HG--
extra : convert_revision : svn%3Ab85faa78-f9eb-468e-a121-7cced6da292c%40213
---
 scrapy/trunk/scrapy/tests/test_decompress.py | 45 +++++++++++++++++
 scrapy/trunk/scrapy/utils/decompressor.py    | 53 ++++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 scrapy/trunk/scrapy/tests/test_decompress.py
 create mode 100644 scrapy/trunk/scrapy/utils/decompressor.py

diff --git a/scrapy/trunk/scrapy/tests/test_decompress.py b/scrapy/trunk/scrapy/tests/test_decompress.py
new file mode 100644
index 000000000..784495b76
--- /dev/null
+++ b/scrapy/trunk/scrapy/tests/test_decompress.py
@@ -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()
diff --git a/scrapy/trunk/scrapy/utils/decompressor.py b/scrapy/trunk/scrapy/utils/decompressor.py
new file mode 100644
index 000000000..60334f3b7
--- /dev/null
+++ b/scrapy/trunk/scrapy/utils/decompressor.py
@@ -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
+