response.set_cookie(), closes #795

datasette/utils/asgi.py

@@ -4,10 +4,15 @@ from mimetypes import guess_type
 from urllib.parse import parse_qs, urlunparse, parse_qsl
 from pathlib import Path
 from html import escape
-from http.cookies import SimpleCookie
+from http.cookies import SimpleCookie, Morsel
 import re
 import aiofiles
 
+# Workaround for adding samesite support to pre 3.8 python
+Morsel._reserved["samesite"] = "SameSite"
+# Thanks, Starlette:
+# https://github.com/encode/starlette/blob/519f575/starlette/responses.py#L17
+
 
 class NotFound(Exception):
     pass
@@ -17,6 +22,9 @@ class Forbidden(Exception):
     pass
 
 
+SAMESITE_VALUES = ("strict", "lax", "none")
+
+
 class Request:
     def __init__(self, scope, receive):
         self.scope = scope
@@ -370,20 +378,24 @@ class Response:
         self.body = body
         self.status = status
         self.headers = headers or {}
+        self._set_cookie_headers = []
         self.content_type = content_type
 
     async def asgi_send(self, send):
         headers = {}
         headers.update(self.headers)
         headers["content-type"] = self.content_type
+        raw_headers = [
+            [key.encode("utf-8"), value.encode("utf-8")]
+            for key, value in headers.items()
+        ]
+        for set_cookie in self._set_cookie_headers:
+            raw_headers.append([b"set-cookie", set_cookie.encode("utf-8")])
         await send(
             {
                 "type": "http.response.start",
                 "status": self.status,
-                "headers": [
-                    [key.encode("utf-8"), value.encode("utf-8")]
-                    for key, value in headers.items()
-                ],
+                "headers": raw_headers,
             }
         )
         body = self.body
@@ -391,6 +403,37 @@ class Response:
             body = body.encode("utf-8")
         await send({"type": "http.response.body", "body": body})
 
+    def set_cookie(
+        self,
+        key,
+        value="",
+        max_age=None,
+        expires=None,
+        path="/",
+        domain=None,
+        secure=False,
+        httponly=False,
+        samesite="lax",
+    ):
+        assert samesite in SAMESITE_VALUES, "samesite should be one of {}".format(
+            SAMESITE_VALUES
+        )
+        cookie = SimpleCookie()
+        cookie[key] = value
+        for prop_name, prop_value in (
+            ("max_age", max_age),
+            ("expires", expires),
+            ("path", path),
+            ("domain", domain),
+            ("samesite", samesite),
+        ):
+            if prop_value is not None:
+                cookie[key][prop_name.replace("_", "-")] = prop_value
+        for prop_name, prop_value in (("secure", secure), ("httponly", httponly)):
+            if prop_value:
+                cookie[key][prop_name] = True
+        self._set_cookie_headers.append(cookie.output(header="").strip())
+
     @classmethod
     def html(cls, body, status=200, headers=None):
         return cls(