Support str and datetime on expires parameter on the set_cookie method (#1908)

Co-authored-by: Hugo Estrada <hugoestrada@cal.berkeley.edu>
Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>
Co-authored-by: Florimond Manca <florimond.manca@protonmail.com>

Author: 4 people

docs/responses.md

@@ -36,7 +36,7 @@ Signature: `Response.set_cookie(key, value, max_age=None, expires=None, path="/"
 * `key` - A string that will be the cookie's key.
 * `value` - A string that will be the cookie's value.
 * `max_age` - An integer that defines the lifetime of the cookie in seconds. A negative integer or a value of `0` will discard the cookie immediately. `Optional`
-* `expires` - An integer that defines the number of seconds until the cookie expires. `Optional`
+* `expires` - Either an integer that defines the number of seconds until the cookie expires, or a datetime. `Optional`
 * `path` - A string that specifies the subset of routes to which the cookie will apply. `Optional`
 * `domain` - A string that specifies the domain for which the cookie is valid. `Optional`
 * `secure` - A bool indicating that the cookie will only be sent to the server if request is made using SSL and the HTTPS protocol. `Optional`

starlette/responses.py

@@ -4,7 +4,8 @@
 import stat
 import sys
 import typing
-from email.utils import formatdate
+from datetime import datetime
+from email.utils import format_datetime, formatdate
 from functools import partial
 from mimetypes import guess_type as mimetypes_guess_type
 from urllib.parse import quote
@@ -105,7 +106,7 @@ def set_cookie(
         key: str,
         value: str = "",
         max_age: typing.Optional[int] = None,
-        expires: typing.Optional[int] = None,
+        expires: typing.Optional[typing.Union[datetime, str, int]] = None,
         path: str = "/",
         domain: typing.Optional[str] = None,
         secure: bool = False,
@@ -117,7 +118,10 @@ def set_cookie(
         if max_age is not None:
             cookie[key]["max-age"] = max_age
         if expires is not None:
-            cookie[key]["expires"] = expires
+            if isinstance(expires, datetime):
+                cookie[key]["expires"] = format_datetime(expires, usegmt=True)
+            else:
+                cookie[key]["expires"] = expires
         if path is not None:
             cookie[key]["path"] = path
         if domain is not None:

tests/test_responses.py

@@ -1,4 +1,7 @@
+import datetime as dt
 import os
+import time
+from http.cookies import SimpleCookie
 
 import anyio
 import pytest
@@ -288,7 +291,11 @@ def test_file_response_with_inline_disposition(tmpdir, test_client_factory):
     assert response.headers["content-disposition"] == expected_disposition
 
 
-def test_set_cookie(test_client_factory):
+def test_set_cookie(test_client_factory, monkeypatch):
+    # Mock time used as a reference for `Expires` by stdlib `SimpleCookie`.
+    mocked_now = dt.datetime(2100, 1, 22, 12, 0, 0, tzinfo=dt.timezone.utc)
+    monkeypatch.setattr(time, "time", lambda: mocked_now.timestamp())
+
     async def app(scope, receive, send):
         response = Response("Hello, world!", media_type="text/plain")
         response.set_cookie(
@@ -307,6 +314,37 @@ async def app(scope, receive, send):
     client = test_client_factory(app)
     response = client.get("/")
     assert response.text == "Hello, world!"
+    assert (
+        response.headers["set-cookie"]
+        == "mycookie=myvalue; Domain=localhost; expires=Fri, 22 Jan 2100 12:00:10 GMT; "
+        "HttpOnly; Max-Age=10; Path=/; SameSite=none; Secure"
+    )
+
+
+@pytest.mark.parametrize(
+    "expires",
+    [
+        pytest.param(
+            dt.datetime(2100, 1, 22, 12, 0, 10, tzinfo=dt.timezone.utc), id="datetime"
+        ),
+        pytest.param("Fri, 22 Jan 2100 12:00:10 GMT", id="str"),
+        pytest.param(10, id="int"),
+    ],
+)
+def test_expires_on_set_cookie(test_client_factory, monkeypatch, expires):
+    # Mock time used as a reference for `Expires` by stdlib `SimpleCookie`.
+    mocked_now = dt.datetime(2100, 1, 22, 12, 0, 0, tzinfo=dt.timezone.utc)
+    monkeypatch.setattr(time, "time", lambda: mocked_now.timestamp())
+
+    async def app(scope, receive, send):
+        response = Response("Hello, world!", media_type="text/plain")
+        response.set_cookie("mycookie", "myvalue", expires=expires)
+        await response(scope, receive, send)
+
+    client = test_client_factory(app)
+    response = client.get("/")
+    cookie: SimpleCookie = SimpleCookie(response.headers.get("set-cookie"))
+    assert cookie["mycookie"]["expires"] == "Fri, 22 Jan 2100 12:00:10 GMT"
 
 
 def test_delete_cookie(test_client_factory):