Repository: adamchainz/apig-wsgi

Issue:
Ability to send binary response with `text/` or `application/json` content type
Problem
--------

`Response.as_apig_response` is hardcoded to not return binary response when the content-type of the response is either `text/*` or `application/json`.

https://github.com/adamchainz/apig-wsgi/blob/b0ce56cbb3ff67586c398f57bc13590bf0961940/apig_wsgi.py#L105-L109

This becomes a problem for use cases like returning a gzip response from the application with the above content types.

Possible Solution
-----------------

**Check if the `Content-Encoding` header has `gzip` in it or not, and return the binary response based on that.**

This fix would keep the hardcoded values, something like this should work:

```python
    def _should_send_binary(self):
        """Determines if binary response should be sent to API Gateway
        """
        non_binary_content_types = ("text/", "application/json")

        content_type = self._get_content_type() or ''
        content_encoding = self._get_content_encoding() or ''

        supports_binary = self.binary_support
        is_binary_content_type = not content_type.startswith(non_binary_content_types)

        if supports_binary and is_binary_content_type:
            return True
        # Content type is non-binary but the content encoding is.
        elif supports_binary and not is_binary_content_type:
            return 'gzip' in content_encoding.lower()

        return False
```

**Allow `Response` class to be extended so that users can handle this logic on their own.**

This can be done by introducing the same `Response._should_send_binary` method which defaults to the current behaviour. The `make_lambda_handler` function can then be allowed to have an argument like `response_class` to facilitate this.

I can create a PR for either of the fixes which seem more appropriate to you.

What do you think about this?


Discussion:

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index c6365ca..3669cc5 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -7,6 +7,7 @@ Pending Release
 .. Insert new release notes below this line
 
 * Update Python support to 3.5-3.7, as 3.4 has reached its end of life.
+* Return binary content for gzipped responses with text or JSON content types.
 
 2.2.0 (2019-04-15)
 ------------------
diff --git a/apig_wsgi.py b/apig_wsgi.py
index aa36d78..be2930f 100644
--- a/apig_wsgi.py
+++ b/apig_wsgi.py
@@ -101,14 +101,7 @@ class Response(object):
 
     def as_apig_response(self):
         response = {"statusCode": self.status_code, "headers": dict(self.headers)}
-
-        content_type = self._get_content_type()
-        should_send_binary = self.binary_support and (
-            content_type is None
-            or not content_type.startswith(("text/", "application/json"))
-        )
-
-        if should_send_binary:
+        if self._should_send_binary():
             response["isBase64Encoded"] = True
             response["body"] = b64encode(self.body.getvalue()).decode("utf-8")
         else:
@@ -116,10 +109,31 @@ class Response(object):
 
         return response
 
+    def _should_send_binary(self):
+        """
+        Determines if binary response should be sent to API Gateway
+        """
+        if not self.binary_support:
+            return False
+
+        content_type = self._get_content_type()
+        non_binary_content_types = ("text/", "application/json")
+        if not content_type.startswith(non_binary_content_types):
+            return True
+
+        content_encoding = self._get_content_encoding()
+        # Content type is non-binary but the content encoding might be.
+        return "gzip" in content_encoding.lower()
+
     def _get_content_type(self):
-        content_type_headers = [
-            v for k, v in self.headers if k.lower() == "content-type"
-        ]
-        if len(content_type_headers):
-            return content_type_headers[-1]
+        return self._get_header("content-type") or ""
+
+    def _get_content_encoding(self):
+        return self._get_header("content-encoding") or ""
+
+    def _get_header(self, header_name):
+        header_name = header_name.lower()
+        matching_headers = [v for k, v in self.headers if k.lower() == header_name]
+        if len(matching_headers):
+            return matching_headers[-1]
         return None


Test patch:
diff --git a/test_apig_wsgi.py b/test_apig_wsgi.py
index 6a921ea..72531fd 100644
--- a/test_apig_wsgi.py
+++ b/test_apig_wsgi.py
@@ -21,6 +21,11 @@ def simple_app():
     yield app
 
 
+parametrize_text_content_type = pytest.mark.parametrize(
+    "text_content_type", ["text/plain", "text/html", "application/json"]
+)
+
+
 def make_event(
     method="GET",
     qs_params=None,
@@ -69,14 +74,15 @@ def test_get_missing_content_type(simple_app):
     assert response == {"statusCode": 200, "headers": {}, "body": "Hello World\n"}
 
 
-def test_get_binary_support_text(simple_app):
+@parametrize_text_content_type
+def test_get_binary_support_text(simple_app, text_content_type):
     simple_app.handler = make_lambda_handler(simple_app, binary_support=True)
+    simple_app.headers = [("Content-Type", text_content_type)]
 
     response = simple_app.handler(make_event(), None)
-
     assert response == {
         "statusCode": 200,
-        "headers": {"Content-Type": "text/plain"},
+        "headers": {"Content-Type": text_content_type},
         "body": "Hello World\n",
     }
 
@@ -96,6 +102,27 @@ def test_get_binary_support_binary(simple_app):
     }
 
 
+@parametrize_text_content_type
+def test_get_binary_support_binary_text_with_gzip_content_encoding(
+    simple_app, text_content_type
+):
+    simple_app.handler = make_lambda_handler(simple_app, binary_support=True)
+    simple_app.headers = [
+        ("Content-Type", text_content_type),
+        ("Content-Encoding", "gzip"),
+    ]
+    simple_app.response = b"\x13\x37"
+
+    response = simple_app.handler(make_event(), None)
+
+    assert response == {
+        "statusCode": 200,
+        "headers": {"Content-Type": text_content_type, "Content-Encoding": "gzip"},
+        "body": b64encode(b"\x13\x37").decode("utf-8"),
+        "isBase64Encoded": True,
+    }
+
+
 def test_get_binary_support_no_content_type(simple_app):
     simple_app.handler = make_lambda_handler(simple_app, binary_support=True)
     simple_app.headers = []