From 0a63a6e586ababc932d57e7751187d4ff2d7ca18 Mon Sep 17 00:00:00 2001 From: oskipa <60695018+oskipa@users.noreply.github.com> Date: Mon, 6 Feb 2023 00:45:30 -0500 Subject: [PATCH] Support `str` and `datetime` on `expires` parameter on the `set_cookie` method (#1908) Co-authored-by: Hugo Estrada Co-authored-by: Marcelo Trylesinski Co-authored-by: Florimond Manca --- docs/responses.md | 2 +- starlette/responses.py | 10 +++++++--- tests/test_responses.py | 40 +++++++++++++++++++++++++++++++++++++++- 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/docs/responses.md b/docs/responses.md index a8c1b78e4..41a87797d 100644 --- a/docs/responses.md +++ b/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` diff --git a/starlette/responses.py b/starlette/responses.py index 3d5b3e43c..453bbb158 100644 --- a/starlette/responses.py +++ b/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: diff --git a/tests/test_responses.py b/tests/test_responses.py index 608842da2..111717384 100644 --- a/tests/test_responses.py +++ b/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):