Merge adc6c2e into 1211a8d

apitools/gen/util.py

@@ -20,11 +20,13 @@
 
 import collections
 import contextlib
+import gzip
 import json
 import keyword
 import logging
 import os
 import re
+import tempfile
 
 import six
 from six.moves import urllib_parse
@@ -365,6 +367,30 @@ def _NormalizeDiscoveryUrls(discovery_url):
     ]
 
 
+def _Gunzip(gzipped_content):
+    """Returns gunzipped content from gzipped contents."""
+    f = tempfile.NamedTemporaryFile(suffix='gz', mode='w+b', delete=False)
+    try:
+        f.write(gzipped_content)
+        f.close()  # force file synchronization
+        with gzip.open(f.name, 'rb') as h:
+            decompressed_content = h.read()
+        return decompressed_content
+    finally:
+        os.unlink(f.name)
+
+
+def _GetURLContent(url):
+    """Download and return the content of URL."""
+    response = urllib_request.urlopen(url)
+    encoding = response.info().get('Content-Encoding')
+    if encoding == 'gzip':
+        content = _Gunzip(response.read())
+    else:
+        content = response.read()
+    return content
+
+
 def FetchDiscoveryDoc(discovery_url, retries=5):
     """Fetch the discovery document at the given url."""
     discovery_urls = _NormalizeDiscoveryUrls(discovery_url)
@@ -373,7 +399,7 @@ def FetchDiscoveryDoc(discovery_url, retries=5):
     for url in discovery_urls:
         for _ in range(retries):
             try:
-                content = urllib_request.urlopen(url).read()
+                content = _GetURLContent(url)
                 if isinstance(content, bytes):
                     content = content.decode('utf8')
                 discovery_doc = json.loads(content)

apitools/gen/util_test.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
 #
 # Copyright 2015 Google Inc.
 #
@@ -14,9 +16,15 @@
 # limitations under the License.
 
 """Tests for util."""
+import codecs
+import gzip
+import os
+import six.moves.urllib.request as urllib_request
+import tempfile
 import unittest2
 
 from apitools.gen import util
+from mock import MagicMock
 
 
 class NormalizeVersionTest(unittest2.TestCase):
@@ -37,3 +45,56 @@ def testKeywords(self):
     def testNormalizeEnumName(self):
         names = util.Names([''])
         self.assertEqual('_0', names.NormalizeEnumName('0'))
+
+
+class MockRequestResponse():
+    """Mocks the behavior of urllib.response."""
+
+    class MockRequestEncoding():
+        def __init__(self, encoding):
+            self.encoding = encoding
+
+        def get(self, _):
+            return self.encoding
+
+    def __init__(self, content, encoding):
+        self.content = content
+        self.encoding = MockRequestResponse.MockRequestEncoding(encoding)
+
+    def info(self):
+        return self.encoding
+
+    def read(self):
+        return self.content
+
+
+def _Gzip(raw_content):
+    """Returns gzipped content from any content."""
+    f = tempfile.NamedTemporaryFile(suffix='gz', mode='wb', delete=False)
+    f.close()
+    try:
+        with gzip.open(f.name, 'wb') as h:
+            h.write(raw_content)
+        with open(f.name, 'rb') as h:
+            return h.read()
+    finally:
+        os.unlink(f.name)
+
+
+class GetURLContentTest(unittest2.TestCase):
+
+    def testUnspecifiedContentEncoding(self):
+        data = 'regular non-gzipped content'
+        urllib_request.urlopen = MagicMock(
+            return_value=MockRequestResponse(data, ''))
+
+        self.assertEqual(data, util._GetURLContent('unused_url_parameter'))
+
+    def testGZippedContent(self):
+        data = u'¿Hola qué tal?'
+        compressed_data = _Gzip(data.encode('utf-8'))
+        urllib_request.urlopen = MagicMock(
+            return_value=MockRequestResponse(compressed_data, 'gzip'))
+
+        self.assertEqual(data, util._GetURLContent(
+            'unused_url_parameter').decode('utf-8'))