Add DebugMiddleware

2018-07-18 13:04:14 +01:00 · 2018-07-18 13:04:14 +01:00 · 9838af83ca
parent bd127f9cc0
commit 9838af83ca
3 changed files with 105 additions and 0 deletions
--- a/README.md
+++ b/README.md
@ -376,6 +376,27 @@ def test_app():

 ---

+## Debugging
+
+You can use Starlette's `DebugMiddleware` to display simple error tracebacks in the browser.
+
+```python
+from starlette.debug import DebugMiddleware
+
+
+class App:
+    def __init__(self, scope):
+        self.scope = scope
+
+    async def __call__(self, receive, send):
+        raise RuntimeError('Something went wrong')
+
+
+app = DebugMiddleware(App)
+```
+
+---
+
 ## Decorators

 The `asgi_application` decorator takes a request/response function and turns
--- a/starlette/debug.py
+++ b/starlette/debug.py
@ -0,0 +1,42 @@
+from starlette.datastructures import Headers
+from starlette.response import HTMLResponse, PlainTextResponse
+import html
+import traceback
+
+
+class DebugMiddleware:
+    def __init__(self, app):
+        self.app = app
+
+    def __call__(self, scope):
+        return _DebugResponder(self.app, scope)
+
+
+class _DebugResponder:
+    def __init__(self, app, scope):
+        self.scope = scope
+        self.asgi_instance = app(scope)
+        self.response_started = False
+
+    async def __call__(self, receive, send):
+        self.raw_send = send
+        try:
+            await self.asgi_instance(receive, self.send)
+        except:
+            if self.response_started:
+                raise
+            headers = Headers(self.scope.get('headers', []))
+            accept = headers.get('accept', '')
+            if 'text/html' in accept:
+                exc_html = html.escape(traceback.format_exc())
+                content = f'<html><body><h1>500 Server Error</h1><pre>{exc_html}</pre></body></html>'
+                response = HTMLResponse(content, status_code=500)
+            else:
+                content = traceback.format_exc()
+                response = PlainTextResponse(content, status_code=500)
+            await response(receive, send)
+
+    async def send(self, message):
+        if message['type'] == 'http.response.start':
+            self.response_started = True
+        await self.raw_send(message)
--- a/tests/test_debug.py
+++ b/tests/test_debug.py
@ -0,0 +1,42 @@
+from starlette import Request, Response, TestClient
+from starlette.debug import DebugMiddleware
+import pytest
+
+
+def test_debug_text():
+    def app(scope):
+        async def asgi(receive, send):
+            raise RuntimeError('Something went wrong')
+        return asgi
+
+    client = TestClient(DebugMiddleware(app))
+    response = client.get("/")
+    assert response.status_code == 500
+    assert response.headers['content-type'].startswith('text/plain')
+    assert 'RuntimeError' in response.text
+
+
+def test_debug_html():
+    def app(scope):
+        async def asgi(receive, send):
+            raise RuntimeError('Something went wrong')
+        return asgi
+
+    client = TestClient(DebugMiddleware(app))
+    response = client.get("/", headers={'Accept': 'text/html, */*'})
+    assert response.status_code == 500
+    assert response.headers['content-type'].startswith('text/html')
+    assert 'RuntimeError' in response.text
+
+
+def test_debug_after_response_sent():
+    def app(scope):
+        async def asgi(receive, send):
+            response = Response(b'', status_code=204)
+            await response(receive, send)
+            raise RuntimeError('Something went wrong')
+        return asgi
+
+    client = TestClient(DebugMiddleware(app))
+    with pytest.raises(RuntimeError):
+        response = client.get("/")