added response compression

Author: fafhrd91

aiohttp/web.py

@@ -393,6 +393,8 @@ def __init__(self, *, status=200, reason=None):
         self._keep_alive = None
         self._chunked = False
         self._chunk_size = None
+        self._compression = False
+        self._compression_force = False
         self._headers = CIMultiDict()
         self._cookies = http.cookies.SimpleCookie()
         self.set_status(status, reason)
@@ -418,6 +420,10 @@ def status(self):
     def chunked(self):
         return self._chunked
 
+    @property
+    def compression(self):
+        return self._compression
+
     @property
     def reason(self):
         return self._reason
@@ -440,6 +446,11 @@ def enable_chunked_encoding(self, chunk_size=None):
         self._chunked = True
         self._chunk_size = chunk_size
 
+    def enable_compression(self, force=False):
+        """Enables response compression with `deflate` encoding."""
+        self._compression = True
+        self._compression_force = force
+
     @property
     def headers(self):
         return self._headers
@@ -568,6 +579,12 @@ def start(self, request):
 
         self._copy_cookies()
 
+        if self._compression:
+            if (self._compression_force or
+                    'deflate' in request.headers.get(
+                        hdrs.ACCEPT_ENCODING, '')):
+                resp_impl.add_compression_filter()
+
         if self._chunked:
             resp_impl.enable_chunked_encoding()
             if self._chunk_size:

tests/test_web_response.py

@@ -1,6 +1,7 @@
 import asyncio
 import unittest
 from unittest import mock
+from aiohttp import hdrs
 from aiohttp.multidict import CIMultiDict
 from aiohttp.web import Request, StreamResponse, Response
 from aiohttp.protocol import RawRequestMessage, HttpVersion11
@@ -138,6 +139,50 @@ def test_chunk_size(self, ResponseImpl):
 
         msg = resp.start(req)
         self.assertTrue(msg.chunked)
+        msg.add_chunking_filter.assert_called_with(8192)
+        self.assertIsNotNone(msg.filter)
+
+    @mock.patch('aiohttp.web.ResponseImpl')
+    def test_compression_no_accept(self, ResponseImpl):
+        req = self.make_request('GET', '/')
+        resp = StreamResponse()
+        self.assertFalse(resp.chunked)
+
+        self.assertFalse(resp.compression)
+        resp.enable_compression()
+        self.assertTrue(resp.compression)
+
+        msg = resp.start(req)
+        self.assertFalse(msg.add_compression_filter.called)
+
+    @mock.patch('aiohttp.web.ResponseImpl')
+    def test_force_compression_no_accept(self, ResponseImpl):
+        req = self.make_request('GET', '/')
+        resp = StreamResponse()
+        self.assertFalse(resp.chunked)
+
+        self.assertFalse(resp.compression)
+        resp.enable_compression(force=True)
+        self.assertTrue(resp.compression)
+
+        msg = resp.start(req)
+        self.assertTrue(msg.add_compression_filter.called)
+        self.assertIsNotNone(msg.filter)
+
+    @mock.patch('aiohttp.web.ResponseImpl')
+    def test_compression(self, ResponseImpl):
+        req = self.make_request(
+            'GET', '/',
+            headers=CIMultiDict({str(hdrs.ACCEPT_ENCODING): 'gzip, deflate'}))
+        resp = StreamResponse()
+        self.assertFalse(resp.chunked)
+
+        self.assertFalse(resp.compression)
+        resp.enable_compression()
+        self.assertTrue(resp.compression)
+
+        msg = resp.start(req)
+        self.assertTrue(msg.add_compression_filter.called)
         self.assertIsNotNone(msg.filter)
 
     def test_write_non_byteish(self):