Repository: encode/starlette

Issue:
Background tasks
* Add a `background=...` keyword argument to all responses. It should be either `None` or a coroutine with no arguments.
* The response `__call__` should run the background coroutine if it exists (after sending the response).
* Provide a `BackgroundTask()` class that wraps a callable and some arguments. This can then be used to pass a function *or* coroutine with args/kwargs to a response.

```python
class BackgroundTask:
    def __init__(self, func, *args, **kwargs):
        self.func = func
        self.args = args
        self.kwargs = kwargs

    async def __call__(self):
        if asyncio.iscoroutine(self.func):
            await self.func(*self.args, **self.kwargs)
        else:
            self.func(*self.args, **self.kwargs)
```

Then we can do...

```
return Response(..., background=BackgroundTask(my_func, a=123)
```

Naming might need some thinking about - not sure that conflicting with asyncio's `Task` is ideal.

This'd be a good feature to get in, since it's a low amount of work, but provides functionality that WSGI is unable to do.

Discussion:
alexbotello: @tomchristie Is this a better documentation example? Thoughts? Changes?

Patch:
diff --git a/docs/background.md b/docs/background.md
new file mode 100644
index 0000000..9963c08
--- /dev/null
+++ b/docs/background.md
@@ -0,0 +1,31 @@
+
+Starlette includes a `BackgroundTask` class for tasks that do not need to be completed before the response is sent.
+
+### Background Task
+
+Signature: `BackgroundTask(func, *args, **kwargs)
+
+```python
+from starlette.applications import Starlette
+from starlette.responses import JSONResponse
+from starlette.background import BackgroundTask
+from email.mime.text import MIMEText
+
+app = Starlette()
+
+@app.route('/user/signup', methods=['POST'])
+async def create_new_user(request):
+    user = await request.json()
+    # Do stuff here
+    task = BackgroundTask(send_welcome_email, user)
+    return JSONResponse(
+        {'msg': 'User successfully created'},
+        background=task
+    )
+
+async def send_welcome_email(info):
+    # Do stuff here
+    message = MIMEText(f"Thank you for registering, {info['name']}")
+    message['To'] = info['email']
+    await send_message(message)
+```
\ No newline at end of file
diff --git a/docs/endpoints.md b/docs/endpoints.md
index b849139..87a5536 100644
--- a/docs/endpoints.md
+++ b/docs/endpoints.md
@@ -1,6 +1,8 @@
 
-Starlette includes an `HTTPEndpoint` class that provides a class-based view pattern which
-handles HTTP method dispatching.
+Starlette includes the classes `HTTPEndpoint` and `WebSocketEndpoint` that provide a class-based view pattern for
+handling HTTP method dispatching and WebSocket sessions.
+
+### HTTPEndpoint
 
 The `HTTPEndpoint` class can be used as an ASGI application:
 
@@ -42,3 +44,104 @@ class User(HTTPEndpoint):
 
 HTTP endpoint classes will respond with "405 Method not allowed" responses for any
 request methods which do not map to a corresponding handler.
+
+### WebSocketEndpoint
+
+The `WebSocketEndpoint` class is an ASGI application that presents a wrapper around
+the functionality of a `WebSocket` instance.
+
+The ASGI connection scope is accessible on the endpoint instance via `.scope` and
+has an attribute `encoding` which may optionally be set, in order to validate the expected websocket data in the `on_receive` method.
+
+The encoding types are:
+
+* `'json'`
+* `'bytes'`
+* `'text'`
+
+There are three overridable methods for handling specific ASGI websocket message types:
+
+* `async def on_connect(websocket, **kwargs)`
+* `async def on_receive(websocket, data)`
+* `async def on_disconnect(websocket, close_code)`
+
+```python
+from starlette.endpoints import WebSocketEndpoint
+
+
+class App(WebSocketEndpoint):
+    encoding = 'bytes'
+
+    async def on_connect(self, websocket, **kwargs):
+        await websocket.accept()
+
+    async def on_receive(self, websocket, data):
+        await websocket.send_bytes(b"Message: " + data)
+
+    async def on_disconnect(self, websocket, close_code):
+        pass
+```
+
+The `WebSocketEndpoint` can also be used with the `Starlette` application class:
+
+```python
+import uvicorn
+from starlette.applications import Starlette
+from starlette.endpoints import WebSocketEndpoint, HTTPEndpoint
+from starlette.responses import HTMLResponse
+
+app = Starlette()
+
+html = """
+<!DOCTYPE html>
+<html>
+    <head>
+        <title>Chat</title>
+    </head>
+    <body>
+        <h1>WebSocket Chat</h1>
+        <form action="" onsubmit="sendMessage(event)">
+            <input type="text" id="messageText" autocomplete="off"/>
+            <button>Send</button>
+        </form>
+        <ul id='messages'>
+        </ul>
+        <script>
+            var ws = new WebSocket("ws://localhost:8000/ws");
+            ws.onmessage = function(event) {
+                var messages = document.getElementById('messages')
+                var message = document.createElement('li')
+                var content = document.createTextNode(event.data)
+                message.appendChild(content)
+                messages.appendChild(message)
+            };
+            function sendMessage(event) {
+                var input = document.getElementById("messageText")
+                ws.send(input.value)
+                input.value = ''
+                event.preventDefault()
+            }
+        </script>
+    </body>
+</html>
+"""
+
+
+@app.route("/")
+class Homepage(HTTPEndpoint):
+    async def get(self, request):
+        return HTMLResponse(html)
+
+
+@app.websocket_route("/ws")
+class Echo(WebSocketEndpoint):
+
+    encoding = "text"
+
+    async def on_receive(self, websocket, data):
+        await websocket.send_text(f"Message text was: {data}")
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+```
diff --git a/starlette/background.py b/starlette/background.py
new file mode 100644
index 0000000..9838735
--- /dev/null
+++ b/starlette/background.py
@@ -0,0 +1,17 @@
+import asyncio
+import functools
+
+
+class BackgroundTask:
+    def __init__(self, func, *args, **kwargs):
+        self.func = func
+        self.args = args
+        self.kwargs = kwargs
+
+    async def __call__(self):
+        if asyncio.iscoroutinefunction(self.func):
+            await asyncio.ensure_future(self.func(*self.args, **self.kwargs))
+        else:
+            fn = functools.partial(self.func, *self.args, **self.kwargs)
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(None, fn)
diff --git a/starlette/debug.py b/starlette/debug.py
index cc6e846..11eca43 100644
--- a/starlette/debug.py
+++ b/starlette/debug.py
@@ -1,10 +1,13 @@
-from starlette.requests import Request
-from starlette.responses import HTMLResponse, PlainTextResponse
 import html
 import traceback
+import typing
+
+from starlette.requests import Request
+from starlette.responses import HTMLResponse, PlainTextResponse, Response
+from starlette.types import Scope, Receive, Send, Message, ASGIApp, ASGIInstance
 
 
-def get_debug_response(request, exc):
+def get_debug_response(request: Request, exc: Exception) -> Response:
     accept = request.headers.get("accept", "")
     if "text/html" in accept:
         exc_html = "".join(traceback.format_tb(exc.__traceback__))
@@ -18,22 +21,22 @@ def get_debug_response(request, exc):
 
 
 class DebugMiddleware:
-    def __init__(self, app):
+    def __init__(self, app: ASGIApp) -> None:
         self.app = app
 
-    def __call__(self, scope):
+    def __call__(self, scope: Scope) -> ASGIInstance:
         if scope["type"] != "http":
             return self.app(scope)
         return _DebugResponder(self.app, scope)
 
 
 class _DebugResponder:
-    def __init__(self, app, scope):
+    def __init__(self, app: ASGIApp, scope: Scope) -> None:
         self.app = app
         self.scope = scope
         self.response_started = False
 
-    async def __call__(self, receive, send):
+    async def __call__(self, receive: Receive, send: Send) -> None:
         self.raw_send = send
         try:
             asgi = self.app(self.scope)
@@ -45,7 +48,7 @@ class _DebugResponder:
                 await response(receive, send)
             raise exc from None
 
-    async def send(self, message):
+    async def send(self, message: Message) -> None:
         if message["type"] == "http.response.start":
             self.response_started = True
         await self.raw_send(message)
diff --git a/starlette/endpoints.py b/starlette/endpoints.py
index 40d0041..608f28a 100644
--- a/starlette/endpoints.py
+++ b/starlette/endpoints.py
@@ -1,8 +1,9 @@
 import asyncio
 import typing
-
+import json
 from starlette.exceptions import HTTPException
 from starlette.requests import Request
+from starlette.websockets import WebSocket
 from starlette.responses import Response, PlainTextResponse
 from starlette.types import Receive, Send, Scope
 
@@ -35,3 +36,67 @@ class HTTPEndpoint:
         if "app" in self.scope:
             raise HTTPException(status_code=405)
         return PlainTextResponse("Method Not Allowed", status_code=405)
+
+
+class WebSocketEndpoint:
+
+    encoding = None  # May be "text", "bytes", or "json".
+
+    def __init__(self, scope: Scope):
+        self.scope = scope
+
+    async def __call__(self, receive: Receive, send: Send):
+        websocket = WebSocket(self.scope, receive=receive, send=send)
+        kwargs = self.scope.get("kwargs", {})
+        await self.on_connect(websocket, **kwargs)
+
+        close_code = None
+
+        try:
+            while True:
+                message = await websocket.receive()
+                if message["type"] == "websocket.receive":
+                    data = await self.decode(websocket, message)
+                    await self.on_receive(websocket, data)
+                elif message["type"] == "websocket.disconnect":
+                    close_code = message.get("code", 1000)
+                    return
+        finally:
+            await self.on_disconnect(websocket, close_code)
+
+    async def decode(self, websocket, message):
+
+        if self.encoding == "text":
+            if "text" not in message:
+                await websocket.close(code=1003)
+                raise RuntimeError("Expected text websocket messages, but got bytes")
+            return message["text"]
+
+        elif self.encoding == "bytes":
+            if "bytes" not in message:
+                await websocket.close(code=1003)
+                raise RuntimeError("Expected bytes websocket messages, but got text")
+            return message["bytes"]
+
+        elif self.encoding == "json":
+            if "bytes" not in message:
+                await websocket.close(code=1003)
+                raise RuntimeError(
+                    "Expected JSON to be transferred as bytes websocket messages, but got text"
+                )
+            return json.loads(message["bytes"].decode("utf-8"))
+
+        assert (
+            self.encoding is None
+        ), f"Unsupported 'encoding' attribute {self.encoding}"
+        return message["text"] if "text" in message else message["bytes"]
+
+    async def on_connect(self, websocket, **kwargs):
+        """Override to handle an incoming websocket connection"""
+        await websocket.accept()
+
+    async def on_receive(self, websocket, data):
+        """Override to handle an incoming websocket message"""
+
+    async def on_disconnect(self, websocket, close_code):
+        """Override to handle a disconnecting websocket"""
diff --git a/starlette/responses.py b/starlette/responses.py
index 369bf92..5ea906b 100644
--- a/starlette/responses.py
+++ b/starlette/responses.py
@@ -1,5 +1,6 @@
 from email.utils import formatdate
 from mimetypes import guess_type
+from starlette.background import BackgroundTask
 from starlette.datastructures import MutableHeaders
 from starlette.types import Receive, Send
 from urllib.parse import quote_plus
@@ -41,11 +42,13 @@ class Response:
         status_code: int = 200,
         headers: dict = None,
         media_type: str = None,
+        background: BackgroundTask = None,
     ) -> None:
         self.body = self.render(content)
         self.status_code = status_code
         if media_type is not None:
             self.media_type = media_type
+        self.background = background
         self.init_headers(headers)
 
     def render(self, content: typing.Any) -> bytes:
@@ -96,6 +99,9 @@ class Response:
         )
         await send({"type": "http.response.body", "body": self.body})
 
+        if self.background is not None:
+            await self.background()
+
 
 class HTMLResponse(Response):
     media_type = "text/html"


Test patch:
diff --git a/tests/test_background.py b/tests/test_background.py
new file mode 100644
index 0000000..1a4c9f5
--- /dev/null
+++ b/tests/test_background.py
@@ -0,0 +1,52 @@
+from starlette.responses import Response
+from starlette.background import BackgroundTask
+from starlette.testclient import TestClient
+import asyncio
+
+
+def test_async_task():
+    async def async_task():
+        count = 0
+        for num in range(3):
+            count += 1
+            await asyncio.sleep(1)
+        return count
+
+    task = BackgroundTask(async_task)
+
+    def app(scope):
+        async def asgi(receive, send):
+            response = Response(
+                "task initiated", media_type="text/plain", background=task
+            )
+            await response(receive, send)
+
+        return asgi
+
+    client = TestClient(app)
+    response = client.get("/")
+    assert response.text == "task initiated"
+
+
+def test_sync_task():
+    def sync_task():
+        num = 500
+        count = 0
+        while count != num:
+            count += 1
+        return count
+
+    task = BackgroundTask(sync_task)
+
+    def app(scope):
+        async def asgi(receive, send):
+            response = Response(
+                "task initiated", media_type="text/plain", background=task
+            )
+            await response(receive, send)
+
+        return asgi
+
+    client = TestClient(app)
+    response = client.get("/")
+    assert response.text == "task initiated"
diff --git a/tests/test_debug.py b/tests/test_debug.py
index a93b277..700c836 100644
--- a/tests/test_debug.py
+++ b/tests/test_debug.py
@@ -46,7 +46,7 @@ def test_debug_after_response_sent():
     app = DebugMiddleware(app)
     client = TestClient(app)
     with pytest.raises(RuntimeError):
-        response = client.get("/")
+        client.get("/")
 
 
 def test_debug_error_during_scope():
diff --git a/tests/test_endpoints.py b/tests/test_endpoints.py
index 4b5a39e..ec93bec 100644
--- a/tests/test_endpoints.py
+++ b/tests/test_endpoints.py
@@ -2,7 +2,7 @@ import pytest
 from starlette.responses import PlainTextResponse
 from starlette.routing import Router, Path
 from starlette.testclient import TestClient
-from starlette.endpoints import HTTPEndpoint
+from starlette.endpoints import HTTPEndpoint, WebSocketEndpoint
 
 
 class Homepage(HTTPEndpoint):
@@ -17,19 +17,109 @@ app = Router(routes=[Path("/", Homepage), Path("/{username}", Homepage)])
 client = TestClient(app)
 
 
-def test_route():
+def test_http_endpoint_route():
     response = client.get("/")
     assert response.status_code == 200
     assert response.text == "Hello, world!"
 
 
-def test_route_kwargs():
+def test_http_endpoint_route_kwargs():
     response = client.get("/tomchristie")
     assert response.status_code == 200
     assert response.text == "Hello, tomchristie!"
 
 
-def test_route_method():
+def test_http_endpoint_route_method():
     response = client.post("/")
     assert response.status_code == 405
     assert response.text == "Method Not Allowed"
+
+
+def test_websocket_endpoint_on_connect():
+    class WebSocketApp(WebSocketEndpoint):
+        async def on_connect(self, websocket, **kwargs):
+            assert websocket["subprotocols"] == ["soap", "wamp"]
+            await websocket.accept(subprotocol="wamp")
+
+    client = TestClient(WebSocketApp)
+    with client.websocket_connect("/ws", subprotocols=["soap", "wamp"]) as websocket:
+        assert websocket.accepted_subprotocol == "wamp"
+
+
+def test_websocket_endpoint_on_receive_bytes():
+    class WebSocketApp(WebSocketEndpoint):
+        encoding = "bytes"
+
+        async def on_receive(self, websocket, data):
+            await websocket.send_bytes(b"Message bytes was: " + data)
+
+    client = TestClient(WebSocketApp)
+    with client.websocket_connect("/ws") as websocket:
+        websocket.send_bytes(b"Hello, world!")
+        _bytes = websocket.receive_bytes()
+        assert _bytes == b"Message bytes was: Hello, world!"
+
+    with pytest.raises(RuntimeError):
+        with client.websocket_connect("/ws") as websocket:
+            websocket.send_text("Hello world")
+
+
+def test_websocket_endpoint_on_receive_json():
+    class WebSocketApp(WebSocketEndpoint):
+        encoding = "json"
+
+        async def on_receive(self, websocket, data):
+            await websocket.send_json({"message": data})
+
+    client = TestClient(WebSocketApp)
+    with client.websocket_connect("/ws") as websocket:
+        websocket.send_json({"hello": "world"})
+        data = websocket.receive_json()
+        assert data == {"message": {"hello": "world"}}
+
+    with pytest.raises(RuntimeError):
+        with client.websocket_connect("/ws") as websocket:
+            websocket.send_text("Hello world")
+
+
+def test_websocket_endpoint_on_receive_text():
+    class WebSocketApp(WebSocketEndpoint):
+        encoding = "text"
+
+        async def on_receive(self, websocket, data):
+            await websocket.send_text(f"Message text was: {data}")
+
+    client = TestClient(WebSocketApp)
+    with client.websocket_connect("/ws") as websocket:
+        websocket.send_text("Hello, world!")
+        _text = websocket.receive_text()
+        assert _text == "Message text was: Hello, world!"
+
+    with pytest.raises(RuntimeError):
+        with client.websocket_connect("/ws") as websocket:
+            websocket.send_bytes(b"Hello world")
+
+
+def test_websocket_endpoint_on_default():
+    class WebSocketApp(WebSocketEndpoint):
+        encoding = None
+
+        async def on_receive(self, websocket, data):
+            await websocket.send_text(f"Message text was: {data}")
+
+    client = TestClient(WebSocketApp)
+    with client.websocket_connect("/ws") as websocket:
+        websocket.send_text("Hello, world!")
+        _text = websocket.receive_text()
+        assert _text == "Message text was: Hello, world!"
+
+
+def test_websocket_endpoint_on_disconnect():
+    class WebSocketApp(WebSocketEndpoint):
+        async def on_disconnect(self, websocket, close_code):
+            assert close_code == 1001
+            await websocket.close(code=close_code)
+
+    client = TestClient(WebSocketApp)
+    with client.websocket_connect("/ws") as websocket:
+        websocket.close(code=1001)