[tor-commits] [stem/master] Decompress gzip responses

[atagar at torproject.org]

commit b02231d2272ffda53a66cd6eb9fa454697c382b5
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Dec 25 13:25:34 2017 -0800

    Decompress gzip responses
    
    Clearly not what we'll have long term, but does the trick to start.
---
 stem/descriptor/collector.py       | 13 ++++++++++++-
 test/integ/descriptor/collector.py | 12 ++++++++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/stem/descriptor/collector.py b/stem/descriptor/collector.py
index ca0e6921..abd666ab 100644
--- a/stem/descriptor/collector.py
+++ b/stem/descriptor/collector.py
@@ -50,6 +50,8 @@ With this you can either download and read directly from CollecTor...
 .. versionadded:: 1.7.0
 """
 
+import gzip
+import io
 import json
 import time
 
@@ -59,7 +61,9 @@ try:
 except ImportError:
   import urllib2 as urllib
 
+import stem.prereq
 import stem.util.enum
+import stem.util.str_tools
 
 Compression = stem.util.enum.Enum('NONE', 'BZ2', 'GZ', 'XZ')
 
@@ -139,7 +143,14 @@ class CollecTor(object):
 
       # TODO: add compression and retry support
 
-      self._cached_index = json.loads(response)
+      if self.compression == Compression.GZ:
+        if stem.prereq.is_python_3():
+          response = gzip.decompress(response)
+        else:
+          # prior to python 3.2 gzip only had GzipFile
+          response = gzip.GzipFile(fileobj = io.BytesIO(response)).read()
+
+      self._cached_index = json.loads(stem.util.str_tools._to_unicode(response))
       self._cached_index_at = time.time()
 
     return self._cached_index
diff --git a/test/integ/descriptor/collector.py b/test/integ/descriptor/collector.py
index 25f5d503..c481b8c7 100644
--- a/test/integ/descriptor/collector.py
+++ b/test/integ/descriptor/collector.py
@@ -12,8 +12,16 @@ from stem.descriptor.collector import CollecTor, Compression
 class TestCollector(unittest.TestCase):
   @test.require.only_run_once
   @test.require.online
-  def test_index(self):
-    collector = CollecTor(compression = Compression.NONE)
+  def test_index_plaintext(self):
+    self._test_index(Compression.NONE)
+
+  @test.require.only_run_once
+  @test.require.online
+  def test_index_gzip(self):
+    self._test_index(Compression.NONE)
+
+  def _test_index(self, compression):
+    collector = CollecTor(compression = compression)
     index = collector.index()
 
     self.assertEqual('https://collector.torproject.org', index['path'])