Merge pull request #117 from ddelange/patch-1

django_brotli/middleware.py

@@ -1,6 +1,7 @@
 # -*- encoding: utf-8 -*-
 # ! python3
 
+import os
 import re
 
 import brotli
@@ -14,10 +15,16 @@
 
 RE_ACCEPT_ENCODING_BROTLI = re.compile(r"\bbr\b")
 MIN_LEN_FOR_RESPONSE_TO_PROCESS = 200
+BROTLI_MODE = getattr(brotli, os.environ.get("BROTLI_MODE", "MODE_GENERIC"))
+BROTLI_QUALITY = int(os.environ.get("BROTLI_QUALITY", 4))
 
 __all__ = ["BrotliMiddleware"]
 
 
+def compress(obj):
+    return brotli.compress(obj, BROTLI_MODE, BROTLI_QUALITY)
+
+
 # noinspection PyClassHasNoInit
 class BrotliMiddleware(MiddlewareMixin):
     """
@@ -56,7 +63,7 @@ def process_response(
             # we won't know the compressed size until we stream it.
             del response["Content-Length"]
         else:
-            compressed_content = brotli.compress(response.content)
+            compressed_content = compress(response.content)
 
             # Return the compressed content only if it's actually shorter.
             if len(compressed_content) >= len(response.content):
@@ -75,7 +82,7 @@ def process_response(
     def compress_stream(self, streaming_content):
         streaming_content = [line.decode("utf-8") for line in list(streaming_content)]
         streaming_content = "".join(streaming_content).encode()
-        streaming_content = map(lambda x: x, [brotli.compress(streaming_content)])
+        streaming_content = map(lambda x: x, [compress(streaming_content)])
 
         return streaming_content