Repository: adamchainz/apig-wsgi

Issue:
Support sending binary body in response
For e.g. assets served through whitenoise

Discussion:
adamchainz: @kthhrv you'll like this

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index 7cbf671..9db579a 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -6,6 +6,9 @@ Pending Release
 
 .. Insert new release notes below this line
 
+* Add ``binary_support`` flag to enable sending binary responses, if enabled on
+  API Gateway.
+
 1.0.0 (2018-03-08)
 ------------------
 
diff --git a/README.rst b/README.rst
index e4ddb10..c5842c0 100644
--- a/README.rst
+++ b/README.rst
@@ -32,3 +32,22 @@ Use **pip**:
     pip install apig-wsgi
 
 Tested on Python 2.7 and Python 3.6.
+
+Usage
+=====
+
+``make_lambda_handler(app, binary_support=False)``
+--------------------------------------------------
+
+``app`` should be a WSGI app, for example from Django's ``wsgi.py`` or Flask's
+``Flask()`` object.
+
+If you want to support sending binary responses, set ``binary_support`` to
+``True`` and make sure you have ``'*/*'`` in the 'binary media types'
+configuration on your Rest API on API Gateway. Note, whilst API Gateway
+supports a list of media types, using '*/*' is the best way to do it, since it
+is used to match the request 'Accept' header as well.
+
+Note that binary responses aren't sent if your response has a 'Content-Type'
+starting 'text/html' or 'application/json' - this is to support sending larger
+text responses.
diff --git a/apig_wsgi.py b/apig_wsgi.py
index daadc9f..0e2fea1 100644
--- a/apig_wsgi.py
+++ b/apig_wsgi.py
@@ -2,6 +2,7 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import sys
+from base64 import b64encode
 from io import BytesIO
 
 import six
@@ -14,14 +15,14 @@ __version__ = '1.0.0'
 __all__ = ('make_lambda_handler',)
 
 
-def make_lambda_handler(wsgi_app):
+def make_lambda_handler(wsgi_app, binary_support=False):
     """
     Turn a WSGI app callable into a Lambda handler function suitable for
     running on API Gateway.
     """
     def handler(event, context):
         environ = get_environ(event)
-        response = Response()
+        response = Response(binary_support=binary_support)
         result = wsgi_app(environ, response.start_response)
         response.consume(result)
         return response.as_apig_response()
@@ -70,10 +71,11 @@ def get_environ(event):
 
 
 class Response(object):
-    def __init__(self):
+    def __init__(self, binary_support):
         self.status_code = '500'
         self.headers = []
         self.body = BytesIO()
+        self.binary_support = binary_support
 
     def start_response(self, status, response_headers, exc_info=None):
         self.status_code = status.split()[0]
@@ -90,8 +92,24 @@ class Response(object):
                 result.close()
 
     def as_apig_response(self):
-        return {
+        response = {
             'statusCode': self.status_code,
             'headers': dict(self.headers),
-            'body': self.body.getvalue().decode('utf-8'),
         }
+
+        content_type = self._get_content_type()
+
+        if self.binary_support and not content_type.startswith(('text/', 'application/json')):
+            response['isBase64Encoded'] = True
+            response['body'] = b64encode(self.body.getvalue()).decode('utf-8')
+        else:
+            response['body'] = self.body.getvalue().decode('utf-8')
+
+        print(response)
+        return response
+
+    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 None


Test patch:
diff --git a/test_apig_wsgi.py b/test_apig_wsgi.py
index a176fa0..a12e90e 100644
--- a/test_apig_wsgi.py
+++ b/test_apig_wsgi.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import, division, print_function, unicode_literals
 
+from base64 import b64encode
 from io import BytesIO
 
 import pytest
@@ -46,6 +47,45 @@ def test_get(simple_app):
     }
 
 
+def test_get_missing_content_type(simple_app):
+    simple_app.headers = []
+
+    response = simple_app.handler(make_event(), None)
+
+    assert response == {
+        'statusCode': '200',
+        'headers': {},
+        'body': 'Hello World\n',
+    }
+
+
+def test_get_binary_support_text(simple_app):
+    simple_app.handler = make_lambda_handler(simple_app, binary_support=True)
+
+    response = simple_app.handler(make_event(), None)
+
+    assert response == {
+        'statusCode': '200',
+        'headers': {'Content-Type': 'text/plain'},
+        'body': 'Hello World\n',
+    }
+
+
+def test_get_binary_support_binary(simple_app):
+    simple_app.handler = make_lambda_handler(simple_app, binary_support=True)
+    simple_app.headers = [('Content-Type', 'application/octet-stream')]
+    simple_app.response = b'\x13\x37'
+
+    response = simple_app.handler(make_event(), None)
+
+    assert response == {
+        'statusCode': '200',
+        'headers': {'Content-Type': 'application/octet-stream'},
+        'body': b64encode(b'\x13\x37').decode('utf-8'),
+        'isBase64Encoded': True,
+    }
+
+
 def test_post(simple_app):
     event = make_event(method='POST', body='The World is Large')