Merge pull request #234 from tobiowo/decode-compressed-response

Decode compressed response option

vcr/config.py

@@ -35,7 +35,8 @@ class VCR(object):
                  before_record_response=None, filter_post_data_parameters=(),
                  match_on=('method', 'scheme', 'host', 'port', 'path', 'query'),
                  before_record=None, inject_cassette=False, serializer='yaml',
-                 cassette_library_dir=None, func_path_generator=None):
+                 cassette_library_dir=None, func_path_generator=None,
+                 decode_compressed_response=False):
         self.serializer = serializer
         self.match_on = match_on
         self.cassette_library_dir = cassette_library_dir
@@ -67,6 +68,7 @@ class VCR(object):
         self.inject_cassette = inject_cassette
         self.path_transformer = path_transformer
         self.func_path_generator = func_path_generator
+        self.decode_compressed_response = decode_compressed_response
         self._custom_patches = tuple(custom_patches)
 
     def _get_serializer(self, serializer_name):
@@ -163,7 +165,12 @@ class VCR(object):
         before_record_response = options.get(
             'before_record_response', self.before_record_response
         )
+        decode_compressed_response = options.get(
+            'decode_compressed_response', self.decode_compressed_response
+        )
         filter_functions = []
+        if decode_compressed_response:
+            filter_functions.append(filters.decode_response)
         if before_record_response:
             if not isinstance(before_record_response, collections.Iterable):
                 before_record_response = (before_record_response,)

vcr/filters.py

@@ -1,6 +1,10 @@
 from six import BytesIO, text_type
 from six.moves.urllib.parse import urlparse, urlencode, urlunparse
+import copy
 import json
+import zlib
+
+from .util import CaseInsensitiveDict
 
 
 def replace_headers(request, replacements):
@@ -120,3 +124,38 @@ def remove_post_data_parameters(request, post_data_parameters_to_remove):
     """
     replacements = [(k, None) for k in post_data_parameters_to_remove]
     return replace_post_data_parameters(request, replacements)
+
+
+def decode_response(response):
+    """
+    If the response is compressed with gzip or deflate:
+      1. decompress the response body
+      2. delete the content-encoding header
+      3. update content-length header to decompressed length
+    """
+    def is_compressed(headers):
+        encoding = headers.get('content-encoding', [])
+        return encoding and encoding[0] in ('gzip', 'deflate')
+
+    def decompress_body(body, encoding):
+        """Returns decompressed body according to encoding using zlib.
+        to (de-)compress gzip format, use wbits = zlib.MAX_WBITS | 16
+        """
+        if encoding == 'gzip':
+            return zlib.decompress(body, zlib.MAX_WBITS | 16)
+        else:  # encoding == 'deflate'
+            return zlib.decompress(body)
+
+    headers = CaseInsensitiveDict(response['headers'])
+    if is_compressed(headers):
+        response = copy.deepcopy(response)
+        encoding = headers['content-encoding'][0]
+        headers['content-encoding'].remove(encoding)
+        if not headers['content-encoding']:
+            del headers['content-encoding']
+
+        new_body = decompress_body(response['body']['string'], encoding)
+        response['body']['string'] = new_body
+        headers['content-length'] = [str(len(new_body))]
+        response['headers'] = dict(headers)
+    return response