Integrated encode/decoder for brotli

3 files changed, 15 insertions, 4 deletions

diff --git a/netlib/encoding.py b/netlib/encoding.py
index da282194..9c8acff7 100644
--- a/netlib/encoding.py
+++ b/netlib/encoding.py
@@ -8,6 +8,7 @@ import collections
 from io import BytesIO
 import gzip
 import zlib
+import brotli
 
 from typing import Union  # noqa
 
@@ -45,7 +46,7 @@ def decode(encoded, encoding, errors='strict'):
             decoded = custom_decode[encoding](encoded)
         except KeyError:
             decoded = codecs.decode(encoded, encoding, errors)
-        if encoding in ("gzip", "deflate"):
+        if encoding in ("gzip", "deflate", "br"):
             _cache = CachedDecode(encoded, encoding, errors, decoded)
         return decoded
     except Exception as e:
@@ -81,7 +82,7 @@ def encode(decoded, encoding, errors='strict'):
             encoded = custom_encode[encoding](decoded)
         except KeyError:
             encoded = codecs.encode(decoded, encoding, errors)
-        if encoding in ("gzip", "deflate"):
+        if encoding in ("gzip", "deflate", "br"):
             _cache = CachedDecode(encoded, encoding, errors, decoded)
         return encoded
     except Exception as e:
@@ -113,6 +114,14 @@ def encode_gzip(content):
     return s.getvalue()
 
 
+def decode_brotli(content):
+    return brotli.decompress(content)
+
+
+def encode_brotli(content):
+    return brotli.compress(content)
+
+
 def decode_deflate(content):
     """
         Returns decompressed data for DEFLATE. Some servers may respond with
@@ -139,11 +148,13 @@ custom_decode = {
     "identity": identity,
     "gzip": decode_gzip,
     "deflate": decode_deflate,
+    "br": decode_brotli,
 }
 custom_encode = {
     "identity": identity,
     "gzip": encode_gzip,
     "deflate": encode_deflate,
+    "br": encode_brotli,
 }
 
 __all__ = ["encode", "decode"]
diff --git a/netlib/http/message.py b/netlib/http/message.py
index be35b8d1..ce92bab1 100644
--- a/netlib/http/message.py
+++ b/netlib/http/message.py
@@ -248,7 +248,7 @@ class Message(basetypes.Serializable):
 
     def encode(self, e):
         """
-        Encodes body with the encoding e, where e is "gzip", "deflate" or "identity".
+        Encodes body with the encoding e, where e is "gzip", "deflate", "identity", or "br".
         Any existing content-encodings are overwritten,
         the content is not decoded beforehand.
 
diff --git a/netlib/http/request.py b/netlib/http/request.py
index 061217a3..d59fead4 100644
--- a/netlib/http/request.py
+++ b/netlib/http/request.py
@@ -337,7 +337,7 @@ class Request(message.Message):
             self.headers["accept-encoding"] = (
                 ', '.join(
                     e
-                    for e in {"gzip", "identity", "deflate"}
+                    for e in {"gzip", "identity", "deflate", "br"}
                     if e in accept_encoding
                 )
             )