Support brotli compression

CHANGES/2518.feature

@@ -0,0 +1 @@
+Add support for ``br`` (brotli) Content-Encoding compression (enabled if ``brotlipy`` is installed).

aiohttp/client_reqrep.py

@@ -15,7 +15,7 @@
 from multidict import CIMultiDict, CIMultiDictProxy, MultiDict, MultiDictProxy
 from yarl import URL
 
-from . import hdrs, helpers, http, multipart, payload
+from . import compression, hdrs, helpers, http, multipart, payload
 from .client_exceptions import (ClientConnectionError, ClientOSError,
                                 ClientResponseError, ContentTypeError,
                                 InvalidURL, ServerFingerprintMismatch)
@@ -152,7 +152,7 @@ class ClientRequest:
 
     DEFAULT_HEADERS = {
         hdrs.ACCEPT: '*/*',
-        hdrs.ACCEPT_ENCODING: 'gzip, deflate',
+        hdrs.ACCEPT_ENCODING: compression.DEFAULT_ACCEPT_ENCODING,
     }
 
     body = b''

aiohttp/compression.py

@@ -1,3 +1,4 @@
+import enum
 import zlib
 
 from .http_exceptions import ContentEncodingError
@@ -9,16 +10,95 @@
     brotli = None
 
 
+if brotli is None:
+    DEFAULT_ACCEPT_ENCODING = 'gzip, deflate'
+else:
+    DEFAULT_ACCEPT_ENCODING = 'gzip, deflate, br'
+
+
+class ContentCoding(enum.Enum):
+    # The content codings that we have support for.
+    #
+    # Additional registered codings are listed at:
+    # https://www.iana.org/assignments/http-parameters/http-parameters.xhtml#content-coding
+    deflate = 'deflate'
+    gzip = 'gzip'
+    identity = 'identity'
+    br = 'br'
+
+    @classmethod
+    def get_from_accept_encoding(cls, accept_encoding):
+        accept_encoding = accept_encoding.lower()
+        for coding in cls:
+            if coding.value in accept_encoding:
+                if coding == cls.br and brotli is None:
+                    continue
+                return coding
+
+    @classmethod
+    def values(cls):
+        _values = getattr(cls, '_values', None)
+        if _values is None:
+            cls._values = _values = frozenset({c.value for c in cls})
+        return _values
+
+
 def get_compressor(encoding):
     if encoding == 'gzip':
-        return zlib.compressobj(wbits=16 + zlib.MAX_WBITS)
+        return ZlibCompressor.gzip()
     elif encoding == 'deflate':
-        return zlib.compressobj(wbits=-zlib.MAX_WBITS)
+        return ZlibCompressor.deflate()
+    elif encoding == 'br':
+        return BrotliCompressor()
+    elif encoding == 'identity':
+        return None
     else:
         raise RuntimeError('Encoding is %s not supported' % encoding)
 
 
+class ZlibCompressor:
+
+    def __init__(self, wbits):
+        self._compress = zlib.compressobj(wbits=wbits)
+        self._finished = False
+
+    @classmethod
+    def gzip(cls):
+        return cls(wbits=16 + zlib.MAX_WBITS)
+
+    @classmethod
+    def deflate(cls):
+        return cls(wbits=-zlib.MAX_WBITS)
+
+    def compress(self, data):
+        return self._compress.compress(data)
+
+    def finish(self):
+        if self._finished:
+            raise RuntimeError('Compressor is finished!')
+        self._finished = True
+        return self._compress.flush()
+
+
+class BrotliCompressor:
+
+    def __init__(self):
+        if brotli is None:  # pragma: no cover
+            raise ContentEncodingError(
+                'Can not decode content-encoding: brotli (br). '
+                'Please install `brotlipy`')
+        self._compress = brotli.Compressor()
+
+    def compress(self, data):
+        return self._compress.compress(data)
+
+    def finish(self):
+        return self._compress.finish()
+
+
 def decompress(encoding, data):
+    if encoding == 'identity':
+        return data
     decompressor = get_decompressor(encoding)
     decompressed = decompressor.decompress(data) + decompressor.flush()
     if not decompressor.eof:
@@ -28,12 +108,12 @@ def decompress(encoding, data):
 
 
 def get_decompressor(encoding):
-    if encoding == 'br':
-        return BrotliDecompressor()
-    elif encoding == 'gzip':
+    if encoding == 'gzip':
         return GzipDecompressor()
     elif encoding == 'deflate':
         return DeflateDecompressor()
+    elif encoding == 'br':
+        return BrotliDecompressor()
     else:
         raise RuntimeError('Encoding %s is not supported' % encoding)
 
@@ -98,6 +178,8 @@ def __init__(self):
         self._eof = None
 
     def decompress(self, chunk):
+        if isinstance(chunk, bytearray):
+            chunk = bytes(chunk)
         return self._decompressor.decompress(chunk)
 
     def flush(self):

aiohttp/http_writer.py

@@ -110,7 +110,7 @@ async def write_eof(self, chunk=b''):
             if chunk:
                 chunk = self._compress.compress(chunk)
 
-            chunk = chunk + self._compress.flush()
+            chunk = chunk + self._compress.finish()
             if chunk and self.chunked:
                 chunk_len = ('%x\r\n' % len(chunk)).encode('ascii')
                 chunk = chunk_len + chunk + b'\r\n0\r\n\r\n'

aiohttp/multipart.py

@@ -9,7 +9,7 @@
 
 from multidict import CIMultiDict
 
-from .compression import decompress, get_compressor
+from .compression import ContentCoding, decompress, get_compressor
 from .hdrs import (CONTENT_DISPOSITION, CONTENT_ENCODING, CONTENT_LENGTH,
                    CONTENT_TRANSFER_ENCODING, CONTENT_TYPE)
 from .helpers import CHAR, TOKEN, parse_mimetype, reify
@@ -386,8 +386,6 @@ def decode(self, data):
 
     def _decode_content(self, data):
         encoding = self.headers[CONTENT_ENCODING].lower()
-        if encoding == 'identity':
-            return data
         return decompress(encoding, data)
 
     def _decode_content_transfer(self, data):
@@ -721,7 +719,7 @@ def append_payload(self, payload):
 
         # compression
         encoding = payload.headers.get(CONTENT_ENCODING, '').lower()
-        if encoding and encoding not in ('deflate', 'gzip', 'identity'):
+        if encoding and encoding not in ContentCoding.values():
             raise RuntimeError('unknown content encoding: {}'.format(encoding))
         if encoding == 'identity':
             encoding = None
@@ -834,7 +832,7 @@ def enable_compression(self, encoding='deflate'):
 
     async def write_eof(self):
         if self._compress is not None:
-            chunk = self._compress.flush()
+            chunk = self._compress.finish()
             if chunk:
                 self._compress = None
                 await self.write(chunk)

aiohttp/web_response.py

@@ -1,6 +1,5 @@
 import collections
 import datetime
-import enum
 import json
 import math
 import time
@@ -11,23 +10,12 @@
 from multidict import CIMultiDict, CIMultiDictProxy
 
 from . import hdrs, payload
-from .compression import get_compressor
+from .compression import ContentCoding, get_compressor
 from .helpers import HeadersMixin, rfc822_formatted_time, sentinel
 from .http import RESPONSES, SERVER_SOFTWARE, HttpVersion10, HttpVersion11
 
 
-__all__ = ('ContentCoding', 'StreamResponse', 'Response', 'json_response')
-
-
-class ContentCoding(enum.Enum):
-    # The content codings that we have support for.
-    #
-    # Additional registered codings are listed at:
-    # https://www.iana.org/assignments/http-parameters/http-parameters.xhtml#content-coding
-    deflate = 'deflate'
-    gzip = 'gzip'
-    identity = 'identity'
-    br = 'br'
+__all__ = ('StreamResponse', 'Response', 'json_response')
 
 
 ############################################################
@@ -284,12 +272,10 @@ def _start_compression(self, request):
         if self._compression_force:
             self._do_start_compression(self._compression_force)
         else:
-            accept_encoding = request.headers.get(
-                hdrs.ACCEPT_ENCODING, '').lower()
-            for coding in ContentCoding:
-                if coding.value in accept_encoding:
-                    self._do_start_compression(coding)
-                    return
+            coding = ContentCoding.get_from_accept_encoding(
+                request.headers.get(hdrs.ACCEPT_ENCODING, ''))
+            if coding:
+                self._do_start_compression(coding)
 
     async def prepare(self, request):
         if self._eof_sent:
@@ -613,7 +599,7 @@ def _do_start_compression(self, coding):
             # compress the whole body
             compressobj = get_compressor(coding.value)
             self._compressed_body = compressobj.compress(self._body) +\
-                compressobj.flush()
+                compressobj.finish()
             self._headers[hdrs.CONTENT_ENCODING] = coding.value
             self._headers[hdrs.CONTENT_LENGTH] = \
                 str(len(self._compressed_body))

docs/client_quickstart.rst

@@ -158,8 +158,8 @@ You can also access the response body as bytes, for non-text requests::
 The ``gzip`` and ``deflate`` transfer-encodings are automatically
 decoded for you.
 
-You can enable ``brotli`` transfer-encodings support,
-just install  `brotlipy <https://github.com/python-hyper/brotlipy>`_.
+``brotli`` transfer-encodings support is enabled if
+`brotlipy <https://github.com/python-hyper/brotlipy>`_ is installed.
 
 JSON Request
 ============

docs/web_reference.rst

@@ -2676,6 +2676,13 @@ Constants
 
       *GZIP compression*
 
+   .. attribute:: br
+
+      *BROTLI compression*
+
+      Supported if  `brotlipy <https://github.com/python-hyper/brotlipy>`_
+      is installed.
+
    .. attribute:: identity
 
       *no compression*

tests/conftest.py

@@ -3,6 +3,8 @@
 import pytest
 from py import path
 
+import aiohttp.compression
+
 
 pytest_plugins = ['aiohttp.pytest_plugin', 'pytester']
 
@@ -15,3 +17,7 @@ def shorttmpdir():
     tmpdir = path.local(tempfile.mkdtemp())
     yield tmpdir
     tmpdir.remove(rec=1)
+
+
+skip_if_no_brotli = pytest.mark.skipif(aiohttp.compression.brotli is None,
+                                       reason='brotlipy not installed')