From 8e323e8d2172bb8bad6d92fd42139dee167210f9 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 09:07:02 -0800 Subject: [PATCH 01/22] Introduce lifespan state --- uvicorn/config.py | 6 +++++- uvicorn/lifespan/__init__.py | 5 +++++ uvicorn/lifespan/off.py | 2 ++ uvicorn/lifespan/on.py | 4 +++- uvicorn/protocols/http/h11_impl.py | 4 ++++ uvicorn/protocols/http/httptools_impl.py | 4 ++++ uvicorn/server.py | 4 ++-- 7 files changed, 25 insertions(+), 4 deletions(-) diff --git a/uvicorn/config.py b/uvicorn/config.py index 0ebc562c1..d429d7ce8 100644 --- a/uvicorn/config.py +++ b/uvicorn/config.py @@ -48,6 +48,8 @@ if TYPE_CHECKING: from asgiref.typing import ASGIApplication + from uvicorn.lifespan.off import LifespanOff + from uvicorn.lifespan.on import LifespanOn HTTPProtocolType = Literal["auto", "h11", "httptools"] WSProtocolType = Literal["auto", "none", "websockets", "wsproto"] @@ -471,7 +473,9 @@ def load(self) -> None: else: self.ws_protocol_class = self.ws - self.lifespan_class = import_from_string(LIFESPAN[self.lifespan]) + self.lifespan_class: Type[Union[LifespanOn, LifespanOff]] = import_from_string( + LIFESPAN[self.lifespan] + ) try: self.loaded_app = import_from_string(self.app) diff --git a/uvicorn/lifespan/__init__.py b/uvicorn/lifespan/__init__.py index e69de29bb..5ceabe12f 100644 --- a/uvicorn/lifespan/__init__.py +++ b/uvicorn/lifespan/__init__.py @@ -0,0 +1,5 @@ +from typing import Union +from uvicorn.lifespan.off import LifespanOff +from uvicorn.lifespan.on import LifespanOn + +Lifespan = Union[LifespanOff, LifespanOn] diff --git a/uvicorn/lifespan/off.py b/uvicorn/lifespan/off.py index 7ec961b5f..fabf4a445 100644 --- a/uvicorn/lifespan/off.py +++ b/uvicorn/lifespan/off.py @@ -1,9 +1,11 @@ +from typing import Any, Dict from uvicorn import Config class LifespanOff: def __init__(self, config: Config) -> None: self.should_exit = False + self.state: Dict[str, Any] = {} async def startup(self) -> None: pass diff --git a/uvicorn/lifespan/on.py b/uvicorn/lifespan/on.py index 0c650aab1..474acfc4d 100644 --- a/uvicorn/lifespan/on.py +++ b/uvicorn/lifespan/on.py @@ -1,7 +1,7 @@ import asyncio import logging from asyncio import Queue -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Any, Dict, Union from uvicorn import Config @@ -42,6 +42,7 @@ def __init__(self, config: Config) -> None: self.startup_failed = False self.shutdown_failed = False self.should_exit = False + self.state: Dict[str, Any] = {} async def startup(self) -> None: self.logger.info("Waiting for application startup.") @@ -82,6 +83,7 @@ async def main(self) -> None: scope: LifespanScope = { "type": "lifespan", "asgi": {"version": self.config.asgi_version, "spec_version": "2.0"}, + "state": self.state, } await app(scope, self.receive, self.send) except BaseException as exc: diff --git a/uvicorn/protocols/http/h11_impl.py b/uvicorn/protocols/http/h11_impl.py index c2764b028..b9f2bf568 100644 --- a/uvicorn/protocols/http/h11_impl.py +++ b/uvicorn/protocols/http/h11_impl.py @@ -40,6 +40,7 @@ HTTPResponseStartEvent, HTTPScope, ) + from uvicorn.lifespan import Lifespan H11Event = Union[ h11.Request, @@ -68,6 +69,7 @@ def __init__( self, config: Config, server_state: ServerState, + lifespan: Lifespan, _loop: Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: @@ -83,6 +85,7 @@ def __init__( self.ws_protocol_class = config.ws_protocol_class self.root_path = config.root_path self.limit_concurrency = config.limit_concurrency + self.lifespan = lifespan # Timeouts self.timeout_keep_alive_task: Optional[asyncio.TimerHandle] = None @@ -223,6 +226,7 @@ def handle_events(self) -> None: "raw_path": raw_path, "query_string": query_string, "headers": self.headers, + "state": self.lifespan.state.copy(), } upgrade = self._get_upgrade() diff --git a/uvicorn/protocols/http/httptools_impl.py b/uvicorn/protocols/http/httptools_impl.py index 734e8945d..15fbf6cc7 100644 --- a/uvicorn/protocols/http/httptools_impl.py +++ b/uvicorn/protocols/http/httptools_impl.py @@ -43,6 +43,7 @@ HTTPResponseStartEvent, HTTPScope, ) + from uvicorn.lifespan import Lifespan HEADER_RE = re.compile(b'[\x00-\x1F\x7F()<>@,;:[]={} \t\\"]') HEADER_VALUE_RE = re.compile(b"[\x00-\x1F\x7F]") @@ -66,6 +67,7 @@ def __init__( self, config: Config, server_state: ServerState, + lifespan: Lifespan, _loop: Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: @@ -81,6 +83,7 @@ def __init__( self.ws_protocol_class = config.ws_protocol_class self.root_path = config.root_path self.limit_concurrency = config.limit_concurrency + self.lifespan = lifespan # Timeouts self.timeout_keep_alive_task: Optional[TimerHandle] = None @@ -237,6 +240,7 @@ def on_message_begin(self) -> None: "scheme": self.scheme, "root_path": self.root_path, "headers": self.headers, + "state": self.lifespan.state.copy(), } # Parser callbacks diff --git a/uvicorn/server.py b/uvicorn/server.py index a3fb31b2b..15e12ee05 100644 --- a/uvicorn/server.py +++ b/uvicorn/server.py @@ -10,7 +10,7 @@ import time from email.utils import formatdate from types import FrameType -from typing import TYPE_CHECKING, List, Optional, Sequence, Set, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Set, Tuple, Union import click @@ -93,7 +93,7 @@ async def startup(self, sockets: Optional[List[socket.socket]] = None) -> None: config = self.config create_protocol = functools.partial( - config.http_protocol_class, config=config, server_state=self.server_state + config.http_protocol_class, config=config, server_state=self.server_state, lifespan=self.lifespan ) loop = asyncio.get_running_loop() From 097d338d559e1d8e681bd4f4e6b19ec9bd2ff11a Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 09:08:26 -0800 Subject: [PATCH 02/22] fix tests and imports --- tests/test_lifespan.py | 2 ++ uvicorn/protocols/http/h11_impl.py | 2 +- uvicorn/protocols/http/httptools_impl.py | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index d49a9dcf9..b0265bcff 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -166,6 +166,7 @@ async def asgi3app(scope, receive, send): assert scope == { "type": "lifespan", "asgi": {"version": "3.0", "spec_version": "2.0"}, + "state": {}, } async def test(): @@ -188,6 +189,7 @@ def asgi2app(scope): assert scope == { "type": "lifespan", "asgi": {"version": "2.0", "spec_version": "2.0"}, + "state": {}, } async def asgi(receive, send): diff --git a/uvicorn/protocols/http/h11_impl.py b/uvicorn/protocols/http/h11_impl.py index b9f2bf568..f64a0269e 100644 --- a/uvicorn/protocols/http/h11_impl.py +++ b/uvicorn/protocols/http/h11_impl.py @@ -69,7 +69,7 @@ def __init__( self, config: Config, server_state: ServerState, - lifespan: Lifespan, + lifespan: "Lifespan", _loop: Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: diff --git a/uvicorn/protocols/http/httptools_impl.py b/uvicorn/protocols/http/httptools_impl.py index 15fbf6cc7..face46885 100644 --- a/uvicorn/protocols/http/httptools_impl.py +++ b/uvicorn/protocols/http/httptools_impl.py @@ -67,7 +67,7 @@ def __init__( self, config: Config, server_state: ServerState, - lifespan: Lifespan, + lifespan: "Lifespan", _loop: Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: From 0cecef370e5be3eabc79a094cada646bef103994 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 09:10:39 -0800 Subject: [PATCH 03/22] Remove unused imports --- uvicorn/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uvicorn/server.py b/uvicorn/server.py index 15e12ee05..57443b1f3 100644 --- a/uvicorn/server.py +++ b/uvicorn/server.py @@ -10,7 +10,7 @@ import time from email.utils import formatdate from types import FrameType -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Set, Tuple, Union +from typing import TYPE_CHECKING, List, Optional, Sequence, Set, Tuple, Union import click From 34216168d366ea4480175f741f5daa88840f9023 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 09:10:52 -0800 Subject: [PATCH 04/22] format --- uvicorn/config.py | 1 + uvicorn/lifespan/__init__.py | 1 + uvicorn/lifespan/off.py | 1 + uvicorn/protocols/http/h11_impl.py | 1 + uvicorn/protocols/http/httptools_impl.py | 1 + uvicorn/server.py | 5 ++++- 6 files changed, 9 insertions(+), 1 deletion(-) diff --git a/uvicorn/config.py b/uvicorn/config.py index d429d7ce8..3fbf59a09 100644 --- a/uvicorn/config.py +++ b/uvicorn/config.py @@ -48,6 +48,7 @@ if TYPE_CHECKING: from asgiref.typing import ASGIApplication + from uvicorn.lifespan.off import LifespanOff from uvicorn.lifespan.on import LifespanOn diff --git a/uvicorn/lifespan/__init__.py b/uvicorn/lifespan/__init__.py index 5ceabe12f..429858a28 100644 --- a/uvicorn/lifespan/__init__.py +++ b/uvicorn/lifespan/__init__.py @@ -1,4 +1,5 @@ from typing import Union + from uvicorn.lifespan.off import LifespanOff from uvicorn.lifespan.on import LifespanOn diff --git a/uvicorn/lifespan/off.py b/uvicorn/lifespan/off.py index fabf4a445..e1516f16a 100644 --- a/uvicorn/lifespan/off.py +++ b/uvicorn/lifespan/off.py @@ -1,4 +1,5 @@ from typing import Any, Dict + from uvicorn import Config diff --git a/uvicorn/protocols/http/h11_impl.py b/uvicorn/protocols/http/h11_impl.py index f64a0269e..0d24d9b80 100644 --- a/uvicorn/protocols/http/h11_impl.py +++ b/uvicorn/protocols/http/h11_impl.py @@ -40,6 +40,7 @@ HTTPResponseStartEvent, HTTPScope, ) + from uvicorn.lifespan import Lifespan H11Event = Union[ diff --git a/uvicorn/protocols/http/httptools_impl.py b/uvicorn/protocols/http/httptools_impl.py index face46885..7f5a8d6df 100644 --- a/uvicorn/protocols/http/httptools_impl.py +++ b/uvicorn/protocols/http/httptools_impl.py @@ -43,6 +43,7 @@ HTTPResponseStartEvent, HTTPScope, ) + from uvicorn.lifespan import Lifespan HEADER_RE = re.compile(b'[\x00-\x1F\x7F()<>@,;:[]={} \t\\"]') diff --git a/uvicorn/server.py b/uvicorn/server.py index 57443b1f3..939f60725 100644 --- a/uvicorn/server.py +++ b/uvicorn/server.py @@ -93,7 +93,10 @@ async def startup(self, sockets: Optional[List[socket.socket]] = None) -> None: config = self.config create_protocol = functools.partial( - config.http_protocol_class, config=config, server_state=self.server_state, lifespan=self.lifespan + config.http_protocol_class, + config=config, + server_state=self.server_state, + lifespan=self.lifespan, ) loop = asyncio.get_running_loop() From 641227af2b23dbaf9729208977fd5248220ff6b3 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 10:00:00 -0800 Subject: [PATCH 05/22] Add support to WebSockets --- uvicorn/protocols/websockets/websockets_impl.py | 4 ++++ uvicorn/protocols/websockets/wsproto_impl.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/uvicorn/protocols/websockets/websockets_impl.py b/uvicorn/protocols/websockets/websockets_impl.py index d7650d179..95eb1f4d8 100644 --- a/uvicorn/protocols/websockets/websockets_impl.py +++ b/uvicorn/protocols/websockets/websockets_impl.py @@ -39,6 +39,7 @@ WebSocketScope, WebSocketSendEvent, ) + from uvicorn.lifespan import Lifespan class Server: @@ -61,6 +62,7 @@ def __init__( self, config: Config, server_state: ServerState, + lifespan: Lifespan, _loop: Optional[asyncio.AbstractEventLoop] = None, ): if not config.loaded: @@ -70,6 +72,7 @@ def __init__( self.app = config.loaded_app self.loop = _loop or asyncio.get_event_loop() self.root_path = config.root_path + self.lifespan = lifespan # Shared server state self.connections = server_state.connections @@ -187,6 +190,7 @@ async def process_request( "query_string": query_string.encode("ascii"), "headers": asgi_headers, "subprotocols": subprotocols, + "state": self.lifespan.state.copy(), } task = self.loop.create_task(self.run_asgi()) task.add_done_callback(self.on_task_complete) diff --git a/uvicorn/protocols/websockets/wsproto_impl.py b/uvicorn/protocols/websockets/wsproto_impl.py index f2677e004..f8ee1bf24 100644 --- a/uvicorn/protocols/websockets/wsproto_impl.py +++ b/uvicorn/protocols/websockets/wsproto_impl.py @@ -32,6 +32,7 @@ WebSocketScope, WebSocketSendEvent, ) + from uvicorn.lifespan import Lifespan WebSocketEvent = typing.Union[ "WebSocketReceiveEvent", @@ -50,6 +51,7 @@ def __init__( self, config: Config, server_state: ServerState, + lifespan: Lifespan, _loop: typing.Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: @@ -60,6 +62,7 @@ def __init__( self.loop = _loop or asyncio.get_event_loop() self.logger = logging.getLogger("uvicorn.error") self.root_path = config.root_path + self.lifespan = lifespan # Shared server state self.connections = server_state.connections @@ -185,6 +188,7 @@ def handle_connect(self, event: events.Request) -> None: "headers": headers, "subprotocols": event.subprotocols, "extensions": None, + "state": self.lifespan.state.copy(), } self.queue.put_nowait({"type": "websocket.connect"}) task = self.loop.create_task(self.run_asgi()) From 942f73ec8fed939ae64bbcbd48e25a333f2b196f Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 10:05:26 -0800 Subject: [PATCH 06/22] fix type annotation --- uvicorn/config.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/uvicorn/config.py b/uvicorn/config.py index 3fbf59a09..c8de0983c 100644 --- a/uvicorn/config.py +++ b/uvicorn/config.py @@ -49,8 +49,7 @@ if TYPE_CHECKING: from asgiref.typing import ASGIApplication - from uvicorn.lifespan.off import LifespanOff - from uvicorn.lifespan.on import LifespanOn + from uvicorn.lifespan import Lifespan HTTPProtocolType = Literal["auto", "h11", "httptools"] WSProtocolType = Literal["auto", "none", "websockets", "wsproto"] @@ -474,7 +473,7 @@ def load(self) -> None: else: self.ws_protocol_class = self.ws - self.lifespan_class: Type[Union[LifespanOn, LifespanOff]] = import_from_string( + self.lifespan_class: "Type[Lifespan]" = import_from_string( LIFESPAN[self.lifespan] ) From 51ae87165fe9a51af0f8c74cb080ead666f4e059 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 10:09:22 -0800 Subject: [PATCH 07/22] add type ignores --- uvicorn/lifespan/on.py | 2 +- uvicorn/protocols/websockets/websockets_impl.py | 1 + uvicorn/protocols/websockets/wsproto_impl.py | 3 ++- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/uvicorn/lifespan/on.py b/uvicorn/lifespan/on.py index 474acfc4d..78cbd9ead 100644 --- a/uvicorn/lifespan/on.py +++ b/uvicorn/lifespan/on.py @@ -83,7 +83,7 @@ async def main(self) -> None: scope: LifespanScope = { "type": "lifespan", "asgi": {"version": self.config.asgi_version, "spec_version": "2.0"}, - "state": self.state, + "state": self.state, # type: ignore[typeddict-item] } await app(scope, self.receive, self.send) except BaseException as exc: diff --git a/uvicorn/protocols/websockets/websockets_impl.py b/uvicorn/protocols/websockets/websockets_impl.py index 95eb1f4d8..44f4f3760 100644 --- a/uvicorn/protocols/websockets/websockets_impl.py +++ b/uvicorn/protocols/websockets/websockets_impl.py @@ -39,6 +39,7 @@ WebSocketScope, WebSocketSendEvent, ) + from uvicorn.lifespan import Lifespan diff --git a/uvicorn/protocols/websockets/wsproto_impl.py b/uvicorn/protocols/websockets/wsproto_impl.py index f8ee1bf24..69ba3632d 100644 --- a/uvicorn/protocols/websockets/wsproto_impl.py +++ b/uvicorn/protocols/websockets/wsproto_impl.py @@ -32,6 +32,7 @@ WebSocketScope, WebSocketSendEvent, ) + from uvicorn.lifespan import Lifespan WebSocketEvent = typing.Union[ @@ -188,7 +189,7 @@ def handle_connect(self, event: events.Request) -> None: "headers": headers, "subprotocols": event.subprotocols, "extensions": None, - "state": self.lifespan.state.copy(), + "state": self.lifespan.state.copy(), # type: ignore[typeddict-item] } self.queue.put_nowait({"type": "websocket.connect"}) task = self.loop.create_task(self.run_asgi()) From c85dac27d4d23d5bd129842e78e2889cf71e3c48 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 13:50:00 -0800 Subject: [PATCH 08/22] add test for http --- tests/protocols/test_http.py | 36 +++++++++++++++++-- .../protocols/websockets/websockets_impl.py | 2 +- uvicorn/protocols/websockets/wsproto_impl.py | 2 +- 3 files changed, 36 insertions(+), 4 deletions(-) diff --git a/tests/protocols/test_http.py b/tests/protocols/test_http.py index c041b69d0..ac84256ff 100644 --- a/tests/protocols/test_http.py +++ b/tests/protocols/test_http.py @@ -2,14 +2,18 @@ import socket import threading import time +from typing import Optional import pytest from tests.response import Response from uvicorn import Server from uvicorn.config import WS_PROTOCOLS, Config +from uvicorn.lifespan.off import LifespanOff from uvicorn.main import ServerState from uvicorn.protocols.http.h11_impl import H11Protocol +from uvicorn.lifespan import Lifespan +from uvicorn.lifespan.on import LifespanOn try: from uvicorn.protocols.http.httptools_impl import HttpToolsProtocol @@ -184,12 +188,13 @@ def add_done_callback(self, callback): pass -def get_connected_protocol(app, protocol_cls, **kwargs): +def get_connected_protocol(app, protocol_cls, lifespan: Optional[Lifespan] = None, **kwargs): loop = MockLoop() transport = MockTransport() config = Config(app=app, **kwargs) + lifespan = lifespan or LifespanOff(config) server_state = ServerState() - protocol = protocol_cls(config=config, server_state=server_state, _loop=loop) + protocol = protocol_cls(config=config, server_state=server_state, lifespan=lifespan, _loop=loop) protocol.connection_made(transport) return protocol @@ -980,3 +985,30 @@ async def app(scope, receive, send): protocol.data_received(SIMPLE_GET_REQUEST) await protocol.loop.run_one() assert b"x-test-header: test value" in protocol.transport.buffer + + +@pytest.mark.anyio +@pytest.mark.parametrize("protocol_cls", HTTP_PROTOCOLS) +async def test_lifespan_state(protocol_cls): + app = Response("Hello, world", media_type="text/plain") + + async def app(scope, receive, send): + assert scope["state"]["a"] == 123 + # modifications to keys are not preserved + scope["state"]["a"] = 234 + # unless of course the value itself is mutated + scope["state"]["b"].append(2) + return await Response("Hi!")(scope, receive, send) + + lifespan = LifespanOn(config=Config(app=app)) + # skip over actually running the lifespan, that is tested + # in the lifespan tests + lifespan.state.update({"a": 123, "b": [1]}) + + protocol = get_connected_protocol(app, protocol_cls, lifespan=lifespan) + protocol.data_received(SIMPLE_GET_REQUEST) + await protocol.loop.run_one() + assert b"HTTP/1.1 200 OK" in protocol.transport.buffer + assert b"Hi!" in protocol.transport.buffer + + assert lifespan.state == {"a": 123, "b": [1, 2]} diff --git a/uvicorn/protocols/websockets/websockets_impl.py b/uvicorn/protocols/websockets/websockets_impl.py index 44f4f3760..6c6825a66 100644 --- a/uvicorn/protocols/websockets/websockets_impl.py +++ b/uvicorn/protocols/websockets/websockets_impl.py @@ -63,7 +63,7 @@ def __init__( self, config: Config, server_state: ServerState, - lifespan: Lifespan, + lifespan: "Lifespan", _loop: Optional[asyncio.AbstractEventLoop] = None, ): if not config.loaded: diff --git a/uvicorn/protocols/websockets/wsproto_impl.py b/uvicorn/protocols/websockets/wsproto_impl.py index 69ba3632d..7464f8904 100644 --- a/uvicorn/protocols/websockets/wsproto_impl.py +++ b/uvicorn/protocols/websockets/wsproto_impl.py @@ -52,7 +52,7 @@ def __init__( self, config: Config, server_state: ServerState, - lifespan: Lifespan, + lifespan: "Lifespan", _loop: typing.Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: From fc0e3d2f35f53e9e9ea98b00c371252c8369241c Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 14:30:42 -0800 Subject: [PATCH 09/22] add test for websockets --- tests/protocols/test_http.py | 32 ++++++++------ tests/protocols/test_websocket.py | 53 ++++++++++++++++++++++++ uvicorn/protocols/http/h11_impl.py | 4 +- uvicorn/protocols/http/httptools_impl.py | 4 +- 4 files changed, 79 insertions(+), 14 deletions(-) diff --git a/tests/protocols/test_http.py b/tests/protocols/test_http.py index ac84256ff..8719307e6 100644 --- a/tests/protocols/test_http.py +++ b/tests/protocols/test_http.py @@ -9,11 +9,11 @@ from tests.response import Response from uvicorn import Server from uvicorn.config import WS_PROTOCOLS, Config +from uvicorn.lifespan import Lifespan from uvicorn.lifespan.off import LifespanOff +from uvicorn.lifespan.on import LifespanOn from uvicorn.main import ServerState from uvicorn.protocols.http.h11_impl import H11Protocol -from uvicorn.lifespan import Lifespan -from uvicorn.lifespan.on import LifespanOn try: from uvicorn.protocols.http.httptools_impl import HttpToolsProtocol @@ -188,13 +188,17 @@ def add_done_callback(self, callback): pass -def get_connected_protocol(app, protocol_cls, lifespan: Optional[Lifespan] = None, **kwargs): +def get_connected_protocol( + app, protocol_cls, lifespan: Optional[Lifespan] = None, **kwargs +): loop = MockLoop() transport = MockTransport() config = Config(app=app, **kwargs) lifespan = lifespan or LifespanOff(config) server_state = ServerState() - protocol = protocol_cls(config=config, server_state=server_state, lifespan=lifespan, _loop=loop) + protocol = protocol_cls( + config=config, server_state=server_state, lifespan=lifespan, _loop=loop + ) protocol.connection_made(transport) return protocol @@ -992,10 +996,13 @@ async def app(scope, receive, send): async def test_lifespan_state(protocol_cls): app = Response("Hello, world", media_type="text/plain") + expected_states = [{"a": 123, "b": [1]}, {"a": 123, "b": [1, 2]}] + async def app(scope, receive, send): - assert scope["state"]["a"] == 123 + expected_state = expected_states.pop(0) + assert scope["state"] == expected_state # modifications to keys are not preserved - scope["state"]["a"] = 234 + scope["state"]["a"] = 456 # unless of course the value itself is mutated scope["state"]["b"].append(2) return await Response("Hi!")(scope, receive, send) @@ -1005,10 +1012,11 @@ async def app(scope, receive, send): # in the lifespan tests lifespan.state.update({"a": 123, "b": [1]}) - protocol = get_connected_protocol(app, protocol_cls, lifespan=lifespan) - protocol.data_received(SIMPLE_GET_REQUEST) - await protocol.loop.run_one() - assert b"HTTP/1.1 200 OK" in protocol.transport.buffer - assert b"Hi!" in protocol.transport.buffer + for _ in range(2): + protocol = get_connected_protocol(app, protocol_cls, lifespan=lifespan) + protocol.data_received(SIMPLE_GET_REQUEST) + await protocol.loop.run_one() + assert b"HTTP/1.1 200 OK" in protocol.transport.buffer + assert b"Hi!" in protocol.transport.buffer - assert lifespan.state == {"a": 123, "b": [1, 2]} + assert not expected_states # consumed diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index c92de84b8..e41c44e6f 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -1065,3 +1065,56 @@ async def open_connection(url): async with run_server(config): headers = await open_connection(f"ws://127.0.0.1:{unused_tcp_port}") assert headers.get_all("Server") == ["uvicorn", "over-ridden", "another-value"] + + +@pytest.mark.anyio +@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS) +@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS) +async def test_lifespan_state(ws_protocol_cls, http_protocol_cls, unused_tcp_port: int): + expected_states = [ + {"a": 123, "b": [1]}, + {"a": 123, "b": [1, 2]}, + ] + + async def lifespan_app(scope, receive, send): + message = await receive() + assert message["type"] == "lifespan.startup" + scope["state"]["a"] = 123 + scope["state"]["b"] = [1] + await send({"type": "lifespan.startup.complete"}) + message = await receive() + assert message["type"] == "lifespan.shutdown" + await send({"type": "lifespan.shutdown.complete"}) + + class App(WebSocketResponse): + async def websocket_connect(self, message): + expected_state = expected_states.pop(0) + assert self.scope["state"] == expected_state + self.scope["state"]["a"] = 456 + self.scope["state"]["b"].append(2) + await self.send({"type": "websocket.accept"}) + + async def open_connection(url): + async with websockets.connect(url) as websocket: + return websocket.open + + async def app_wrapper(scope, receive, send): + if scope["type"] == "lifespan": + return await lifespan_app(scope, receive, send) + else: + return await App(scope, receive, send) + + config = Config( + app=app_wrapper, + ws=ws_protocol_cls, + http=http_protocol_cls, + lifespan="on", + port=unused_tcp_port, + ) + async with run_server(config): + is_open = await open_connection(f"ws://127.0.0.1:{unused_tcp_port}") + assert is_open + is_open = await open_connection(f"ws://127.0.0.1:{unused_tcp_port}") + assert is_open + + assert not expected_states # consumed diff --git a/uvicorn/protocols/http/h11_impl.py b/uvicorn/protocols/http/h11_impl.py index 0d24d9b80..37ede9172 100644 --- a/uvicorn/protocols/http/h11_impl.py +++ b/uvicorn/protocols/http/h11_impl.py @@ -289,7 +289,9 @@ def handle_websocket_upgrade(self, event: H11Event) -> None: output += [name, b": ", value, b"\r\n"] output.append(b"\r\n") protocol = self.ws_protocol_class( # type: ignore[call-arg, misc] - config=self.config, server_state=self.server_state + config=self.config, + server_state=self.server_state, + lifespan=self.lifespan, ) protocol.connection_made(self.transport) protocol.data_received(b"".join(output)) diff --git a/uvicorn/protocols/http/httptools_impl.py b/uvicorn/protocols/http/httptools_impl.py index 7f5a8d6df..02d8cfe6b 100644 --- a/uvicorn/protocols/http/httptools_impl.py +++ b/uvicorn/protocols/http/httptools_impl.py @@ -205,7 +205,9 @@ def handle_websocket_upgrade(self) -> None: output += [name, b": ", value, b"\r\n"] output.append(b"\r\n") protocol = self.ws_protocol_class( # type: ignore[call-arg, misc] - config=self.config, server_state=self.server_state + config=self.config, + server_state=self.server_state, + lifespan=self.lifespan, ) protocol.connection_made(self.transport) protocol.data_received(b"".join(output)) From d35d1c61e45e27e49a14a9064b74ed31c0c997a8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 14:32:28 -0800 Subject: [PATCH 10/22] Add test for lifespan itself --- tests/test_lifespan.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index b0265bcff..e481ce7fe 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -247,3 +247,27 @@ async def test(): assert "the lifespan event failed" in error_messages.pop(0) assert "Application shutdown failed. Exiting." in error_messages.pop(0) loop.close() + + +def test_lifespan_state(): + + async def app(scope, receive, send): + message = await receive() + assert message["type"] == "lifespan.startup" + await send({"type": "lifespan.startup.complete"}) + scope["state"]["foo"] = 123 + message = await receive() + assert message["type"] == "lifespan.shutdown" + await send({"type": "lifespan.shutdown.complete"}) + + async def test(): + config = Config(app=app, lifespan="on") + lifespan = LifespanOn(config) + + await lifespan.startup() + assert lifespan.state == {"foo": 123} + await lifespan.shutdown() + + loop = asyncio.new_event_loop() + loop.run_until_complete(test()) + loop.close() From ad4bf56e1749c336b0f1312c0f6cd0ffef63e6b4 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 16:44:57 -0800 Subject: [PATCH 11/22] lint --- tests/test_lifespan.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index e481ce7fe..a9cb73e3a 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -250,7 +250,6 @@ async def test(): def test_lifespan_state(): - async def app(scope, receive, send): message = await receive() assert message["type"] == "lifespan.startup" From 39b0e4ca3a924cb2c8e23172f323aa99c5fc433d Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 19:57:45 -0500 Subject: [PATCH 12/22] fix tests --- tests/test_auto_detection.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_auto_detection.py b/tests/test_auto_detection.py index 9d596df05..eb8ba30bd 100644 --- a/tests/test_auto_detection.py +++ b/tests/test_auto_detection.py @@ -3,6 +3,7 @@ import pytest from uvicorn.config import Config +from uvicorn.lifespan.off import LifespanOff from uvicorn.loops.auto import auto_loop_setup from uvicorn.main import ServerState from uvicorn.protocols.http.auto import AutoHTTPProtocol @@ -45,7 +46,9 @@ def test_loop_auto(): async def test_http_auto(): config = Config(app=app) server_state = ServerState() - protocol = AutoHTTPProtocol(config=config, server_state=server_state) + protocol = AutoHTTPProtocol( + config=config, server_state=server_state, lifespan=LifespanOff(config=config) + ) expected_http = "H11Protocol" if httptools is None else "HttpToolsProtocol" assert type(protocol).__name__ == expected_http @@ -54,6 +57,8 @@ async def test_http_auto(): async def test_websocket_auto(): config = Config(app=app) server_state = ServerState() - protocol = AutoWebSocketsProtocol(config=config, server_state=server_state) + protocol = AutoWebSocketsProtocol( + config=config, server_state=server_state, lifespan=LifespanOff(config=config) + ) expected_websockets = "WSProtocol" if websockets is None else "WebSocketProtocol" assert type(protocol).__name__ == expected_websockets From caa22b395acdac3d48cdf896af0c43f47c65e1b0 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 20:01:26 -0500 Subject: [PATCH 13/22] move type ignore? --- uvicorn/protocols/websockets/wsproto_impl.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uvicorn/protocols/websockets/wsproto_impl.py b/uvicorn/protocols/websockets/wsproto_impl.py index 7464f8904..8e262cc76 100644 --- a/uvicorn/protocols/websockets/wsproto_impl.py +++ b/uvicorn/protocols/websockets/wsproto_impl.py @@ -175,7 +175,7 @@ def handle_connect(self, event: events.Request) -> None: headers = [(b"host", event.host.encode())] headers += [(key.lower(), value) for key, value in event.extra_headers] raw_path, _, query_string = event.target.partition("?") - self.scope: "WebSocketScope" = { + self.scope: "WebSocketScope" = { # type: ignore[typeddict-item] "type": "websocket", "asgi": {"version": self.config.asgi_version, "spec_version": "2.3"}, "http_version": "1.1", @@ -189,7 +189,7 @@ def handle_connect(self, event: events.Request) -> None: "headers": headers, "subprotocols": event.subprotocols, "extensions": None, - "state": self.lifespan.state.copy(), # type: ignore[typeddict-item] + "state": self.lifespan.state.copy(), } self.queue.put_nowait({"type": "websocket.connect"}) task = self.loop.create_task(self.run_asgi()) From a6d0af803876780ad2b541b08e81d0dcbbde61a4 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Dec 2022 20:04:37 -0500 Subject: [PATCH 14/22] move another type ignore --- uvicorn/lifespan/on.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uvicorn/lifespan/on.py b/uvicorn/lifespan/on.py index 78cbd9ead..37e935f01 100644 --- a/uvicorn/lifespan/on.py +++ b/uvicorn/lifespan/on.py @@ -80,10 +80,10 @@ async def shutdown(self) -> None: async def main(self) -> None: try: app = self.config.loaded_app - scope: LifespanScope = { + scope: LifespanScope = { # type: ignore[typeddict-item] "type": "lifespan", "asgi": {"version": self.config.asgi_version, "spec_version": "2.0"}, - "state": self.state, # type: ignore[typeddict-item] + "state": self.state, } await app(scope, self.receive, self.send) except BaseException as exc: From c199bc7eaed5b0747358bdcfcdd4466b257cf9cd Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 28 Dec 2022 05:58:30 -0500 Subject: [PATCH 15/22] pass around dicts and minimize changes --- tests/protocols/test_http.py | 13 +++++++---- tests/test_auto_detection.py | 7 ++---- uvicorn/config.py | 6 +---- uvicorn/lifespan/__init__.py | 6 ----- uvicorn/protocols/http/h11_impl.py | 21 +++++++++++++----- uvicorn/protocols/http/httptools_impl.py | 22 ++++++++++++++----- uvicorn/protocols/websockets/auto.py | 2 +- .../protocols/websockets/websockets_impl.py | 20 ++++++++++++----- uvicorn/protocols/websockets/wsproto_impl.py | 8 +++---- uvicorn/server.py | 17 ++++++++------ 10 files changed, 71 insertions(+), 51 deletions(-) diff --git a/tests/protocols/test_http.py b/tests/protocols/test_http.py index 8719307e6..dd31a5bb4 100644 --- a/tests/protocols/test_http.py +++ b/tests/protocols/test_http.py @@ -2,14 +2,13 @@ import socket import threading import time -from typing import Optional +from typing import Optional, Union import pytest from tests.response import Response from uvicorn import Server from uvicorn.config import WS_PROTOCOLS, Config -from uvicorn.lifespan import Lifespan from uvicorn.lifespan.off import LifespanOff from uvicorn.lifespan.on import LifespanOn from uvicorn.main import ServerState @@ -189,7 +188,10 @@ def add_done_callback(self, callback): def get_connected_protocol( - app, protocol_cls, lifespan: Optional[Lifespan] = None, **kwargs + app, + protocol_cls, + lifespan: Optional[Union[LifespanOff, LifespanOn]] = None, + **kwargs, ): loop = MockLoop() transport = MockTransport() @@ -197,7 +199,10 @@ def get_connected_protocol( lifespan = lifespan or LifespanOff(config) server_state = ServerState() protocol = protocol_cls( - config=config, server_state=server_state, lifespan=lifespan, _loop=loop + config=config, + server_state=server_state, + app_state=lifespan.state.copy(), + _loop=loop, ) protocol.connection_made(transport) return protocol diff --git a/tests/test_auto_detection.py b/tests/test_auto_detection.py index eb8ba30bd..2cd2f70c4 100644 --- a/tests/test_auto_detection.py +++ b/tests/test_auto_detection.py @@ -3,7 +3,6 @@ import pytest from uvicorn.config import Config -from uvicorn.lifespan.off import LifespanOff from uvicorn.loops.auto import auto_loop_setup from uvicorn.main import ServerState from uvicorn.protocols.http.auto import AutoHTTPProtocol @@ -46,9 +45,7 @@ def test_loop_auto(): async def test_http_auto(): config = Config(app=app) server_state = ServerState() - protocol = AutoHTTPProtocol( - config=config, server_state=server_state, lifespan=LifespanOff(config=config) - ) + protocol = AutoHTTPProtocol(config=config, server_state=server_state, app_state={}) expected_http = "H11Protocol" if httptools is None else "HttpToolsProtocol" assert type(protocol).__name__ == expected_http @@ -58,7 +55,7 @@ async def test_websocket_auto(): config = Config(app=app) server_state = ServerState() protocol = AutoWebSocketsProtocol( - config=config, server_state=server_state, lifespan=LifespanOff(config=config) + config=config, server_state=server_state, app_state={} ) expected_websockets = "WSProtocol" if websockets is None else "WebSocketProtocol" assert type(protocol).__name__ == expected_websockets diff --git a/uvicorn/config.py b/uvicorn/config.py index c8de0983c..0ebc562c1 100644 --- a/uvicorn/config.py +++ b/uvicorn/config.py @@ -49,8 +49,6 @@ if TYPE_CHECKING: from asgiref.typing import ASGIApplication - from uvicorn.lifespan import Lifespan - HTTPProtocolType = Literal["auto", "h11", "httptools"] WSProtocolType = Literal["auto", "none", "websockets", "wsproto"] LifespanType = Literal["auto", "on", "off"] @@ -473,9 +471,7 @@ def load(self) -> None: else: self.ws_protocol_class = self.ws - self.lifespan_class: "Type[Lifespan]" = import_from_string( - LIFESPAN[self.lifespan] - ) + self.lifespan_class = import_from_string(LIFESPAN[self.lifespan]) try: self.loaded_app = import_from_string(self.app) diff --git a/uvicorn/lifespan/__init__.py b/uvicorn/lifespan/__init__.py index 429858a28..e69de29bb 100644 --- a/uvicorn/lifespan/__init__.py +++ b/uvicorn/lifespan/__init__.py @@ -1,6 +0,0 @@ -from typing import Union - -from uvicorn.lifespan.off import LifespanOff -from uvicorn.lifespan.on import LifespanOn - -Lifespan = Union[LifespanOff, LifespanOn] diff --git a/uvicorn/protocols/http/h11_impl.py b/uvicorn/protocols/http/h11_impl.py index 37ede9172..46afbbd25 100644 --- a/uvicorn/protocols/http/h11_impl.py +++ b/uvicorn/protocols/http/h11_impl.py @@ -2,7 +2,17 @@ import http import logging import sys -from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Union, + cast, +) from urllib.parse import unquote import h11 @@ -41,7 +51,6 @@ HTTPScope, ) - from uvicorn.lifespan import Lifespan H11Event = Union[ h11.Request, @@ -70,7 +79,7 @@ def __init__( self, config: Config, server_state: ServerState, - lifespan: "Lifespan", + app_state: Dict[str, Any], _loop: Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: @@ -86,7 +95,7 @@ def __init__( self.ws_protocol_class = config.ws_protocol_class self.root_path = config.root_path self.limit_concurrency = config.limit_concurrency - self.lifespan = lifespan + self.app_state = app_state # Timeouts self.timeout_keep_alive_task: Optional[asyncio.TimerHandle] = None @@ -227,7 +236,7 @@ def handle_events(self) -> None: "raw_path": raw_path, "query_string": query_string, "headers": self.headers, - "state": self.lifespan.state.copy(), + "state": self.app_state, } upgrade = self._get_upgrade() @@ -291,7 +300,7 @@ def handle_websocket_upgrade(self, event: H11Event) -> None: protocol = self.ws_protocol_class( # type: ignore[call-arg, misc] config=self.config, server_state=self.server_state, - lifespan=self.lifespan, + app_state=self.app_state, ) protocol.connection_made(self.transport) protocol.data_received(b"".join(output)) diff --git a/uvicorn/protocols/http/httptools_impl.py b/uvicorn/protocols/http/httptools_impl.py index 02d8cfe6b..d7e6c33f0 100644 --- a/uvicorn/protocols/http/httptools_impl.py +++ b/uvicorn/protocols/http/httptools_impl.py @@ -6,7 +6,18 @@ import urllib from asyncio.events import TimerHandle from collections import deque -from typing import TYPE_CHECKING, Callable, Deque, List, Optional, Tuple, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Deque, + Dict, + List, + Optional, + Tuple, + Union, + cast, +) import httptools @@ -44,7 +55,6 @@ HTTPScope, ) - from uvicorn.lifespan import Lifespan HEADER_RE = re.compile(b'[\x00-\x1F\x7F()<>@,;:[]={} \t\\"]') HEADER_VALUE_RE = re.compile(b"[\x00-\x1F\x7F]") @@ -68,7 +78,7 @@ def __init__( self, config: Config, server_state: ServerState, - lifespan: "Lifespan", + app_state: Dict[str, Any], _loop: Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: @@ -84,7 +94,7 @@ def __init__( self.ws_protocol_class = config.ws_protocol_class self.root_path = config.root_path self.limit_concurrency = config.limit_concurrency - self.lifespan = lifespan + self.app_state = app_state # Timeouts self.timeout_keep_alive_task: Optional[TimerHandle] = None @@ -207,7 +217,7 @@ def handle_websocket_upgrade(self) -> None: protocol = self.ws_protocol_class( # type: ignore[call-arg, misc] config=self.config, server_state=self.server_state, - lifespan=self.lifespan, + app_state=self.app_state, ) protocol.connection_made(self.transport) protocol.data_received(b"".join(output)) @@ -243,7 +253,7 @@ def on_message_begin(self) -> None: "scheme": self.scheme, "root_path": self.root_path, "headers": self.headers, - "state": self.lifespan.state.copy(), + "state": self.app_state, } # Parser callbacks diff --git a/uvicorn/protocols/websockets/auto.py b/uvicorn/protocols/websockets/auto.py index 0dfba3bdb..368b98242 100644 --- a/uvicorn/protocols/websockets/auto.py +++ b/uvicorn/protocols/websockets/auto.py @@ -1,7 +1,7 @@ import asyncio import typing -AutoWebSocketsProtocol: typing.Optional[typing.Type[asyncio.Protocol]] +AutoWebSocketsProtocol: typing.Optional[typing.Callable[..., asyncio.Protocol]] try: import websockets # noqa except ImportError: # pragma: no cover diff --git a/uvicorn/protocols/websockets/websockets_impl.py b/uvicorn/protocols/websockets/websockets_impl.py index 6c6825a66..88a1b568d 100644 --- a/uvicorn/protocols/websockets/websockets_impl.py +++ b/uvicorn/protocols/websockets/websockets_impl.py @@ -2,7 +2,17 @@ import http import logging import sys -from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Optional, + Sequence, + Tuple, + Union, + cast, +) from urllib.parse import unquote import websockets @@ -40,8 +50,6 @@ WebSocketSendEvent, ) - from uvicorn.lifespan import Lifespan - class Server: closing = False @@ -63,7 +71,7 @@ def __init__( self, config: Config, server_state: ServerState, - lifespan: "Lifespan", + app_state: Dict[str, Any], _loop: Optional[asyncio.AbstractEventLoop] = None, ): if not config.loaded: @@ -73,7 +81,7 @@ def __init__( self.app = config.loaded_app self.loop = _loop or asyncio.get_event_loop() self.root_path = config.root_path - self.lifespan = lifespan + self.app_state = app_state # Shared server state self.connections = server_state.connections @@ -191,7 +199,7 @@ async def process_request( "query_string": query_string.encode("ascii"), "headers": asgi_headers, "subprotocols": subprotocols, - "state": self.lifespan.state.copy(), + "state": self.app_state, } task = self.loop.create_task(self.run_asgi()) task.add_done_callback(self.on_task_complete) diff --git a/uvicorn/protocols/websockets/wsproto_impl.py b/uvicorn/protocols/websockets/wsproto_impl.py index 8e262cc76..00e15741f 100644 --- a/uvicorn/protocols/websockets/wsproto_impl.py +++ b/uvicorn/protocols/websockets/wsproto_impl.py @@ -33,8 +33,6 @@ WebSocketSendEvent, ) - from uvicorn.lifespan import Lifespan - WebSocketEvent = typing.Union[ "WebSocketReceiveEvent", "WebSocketDisconnectEvent", @@ -52,7 +50,7 @@ def __init__( self, config: Config, server_state: ServerState, - lifespan: "Lifespan", + app_state: typing.Dict[str, typing.Any], _loop: typing.Optional[asyncio.AbstractEventLoop] = None, ) -> None: if not config.loaded: @@ -63,7 +61,7 @@ def __init__( self.loop = _loop or asyncio.get_event_loop() self.logger = logging.getLogger("uvicorn.error") self.root_path = config.root_path - self.lifespan = lifespan + self.app_state = app_state # Shared server state self.connections = server_state.connections @@ -189,7 +187,7 @@ def handle_connect(self, event: events.Request) -> None: "headers": headers, "subprotocols": event.subprotocols, "extensions": None, - "state": self.lifespan.state.copy(), + "state": self.app_state, } self.queue.put_nowait({"type": "websocket.connect"}) task = self.loop.create_task(self.run_asgi()) diff --git a/uvicorn/server.py b/uvicorn/server.py index 939f60725..59866f1e3 100644 --- a/uvicorn/server.py +++ b/uvicorn/server.py @@ -1,5 +1,4 @@ import asyncio -import functools import logging import os import platform @@ -92,12 +91,16 @@ async def startup(self, sockets: Optional[List[socket.socket]] = None) -> None: config = self.config - create_protocol = functools.partial( - config.http_protocol_class, - config=config, - server_state=self.server_state, - lifespan=self.lifespan, - ) + def create_protocol( + _loop: Optional[asyncio.AbstractEventLoop] = None, + ) -> asyncio.Protocol: + return config.http_protocol_class( + config=config, + server_state=self.server_state, + app_state=self.lifespan.state.copy(), + _loop=_loop, + ) + loop = asyncio.get_running_loop() listeners: Sequence[socket.SocketType] From 55419e8f83e6ec0d335848b8992d4a68d2470831 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 28 Dec 2022 06:03:08 -0500 Subject: [PATCH 16/22] move to extensions["state"] --- tests/test_lifespan.py | 4 ++-- uvicorn/lifespan/on.py | 2 +- uvicorn/server.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index a9cb73e3a..cbc1a3ba2 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -166,7 +166,7 @@ async def asgi3app(scope, receive, send): assert scope == { "type": "lifespan", "asgi": {"version": "3.0", "spec_version": "2.0"}, - "state": {}, + "extensions": {"state": {}}, } async def test(): @@ -189,7 +189,7 @@ def asgi2app(scope): assert scope == { "type": "lifespan", "asgi": {"version": "2.0", "spec_version": "2.0"}, - "state": {}, + "extensions": {"state": {}}, } async def asgi(receive, send): diff --git a/uvicorn/lifespan/on.py b/uvicorn/lifespan/on.py index 37e935f01..7bcc3725a 100644 --- a/uvicorn/lifespan/on.py +++ b/uvicorn/lifespan/on.py @@ -83,7 +83,7 @@ async def main(self) -> None: scope: LifespanScope = { # type: ignore[typeddict-item] "type": "lifespan", "asgi": {"version": self.config.asgi_version, "spec_version": "2.0"}, - "state": self.state, + "extensions": {"state": self.state}, } await app(scope, self.receive, self.send) except BaseException as exc: diff --git a/uvicorn/server.py b/uvicorn/server.py index 59866f1e3..2d6c02ff0 100644 --- a/uvicorn/server.py +++ b/uvicorn/server.py @@ -94,7 +94,7 @@ async def startup(self, sockets: Optional[List[socket.socket]] = None) -> None: def create_protocol( _loop: Optional[asyncio.AbstractEventLoop] = None, ) -> asyncio.Protocol: - return config.http_protocol_class( + return config.http_protocol_class( # type: ignore[call-arg] config=config, server_state=self.server_state, app_state=self.lifespan.state.copy(), From dac71df8d18cf55fe387b8b4d995d3488670e9d1 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 28 Dec 2022 06:07:49 -0500 Subject: [PATCH 17/22] fix tests --- tests/protocols/test_websocket.py | 4 ++-- tests/test_lifespan.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index e41c44e6f..20aea5951 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -1079,8 +1079,8 @@ async def test_lifespan_state(ws_protocol_cls, http_protocol_cls, unused_tcp_por async def lifespan_app(scope, receive, send): message = await receive() assert message["type"] == "lifespan.startup" - scope["state"]["a"] = 123 - scope["state"]["b"] = [1] + scope["extensions"]["state"]["a"] = 123 + scope["extensions"]["state"]["b"] = [1] await send({"type": "lifespan.startup.complete"}) message = await receive() assert message["type"] == "lifespan.shutdown" diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index cbc1a3ba2..44a65c8b6 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -254,7 +254,7 @@ async def app(scope, receive, send): message = await receive() assert message["type"] == "lifespan.startup" await send({"type": "lifespan.startup.complete"}) - scope["state"]["foo"] = 123 + scope["extensions"]["state"]["foo"] = 123 message = await receive() assert message["type"] == "lifespan.shutdown" await send({"type": "lifespan.shutdown.complete"}) From 8b0589787d0113c1f2fed42c17dfe30bf0852ff4 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Sat, 4 Mar 2023 12:51:12 +0100 Subject: [PATCH 18/22] move state to top level namespace --- tests/protocols/test_websocket.py | 4 ++-- tests/test_lifespan.py | 6 +++--- uvicorn/lifespan/on.py | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index 20aea5951..e41c44e6f 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -1079,8 +1079,8 @@ async def test_lifespan_state(ws_protocol_cls, http_protocol_cls, unused_tcp_por async def lifespan_app(scope, receive, send): message = await receive() assert message["type"] == "lifespan.startup" - scope["extensions"]["state"]["a"] = 123 - scope["extensions"]["state"]["b"] = [1] + scope["state"]["a"] = 123 + scope["state"]["b"] = [1] await send({"type": "lifespan.startup.complete"}) message = await receive() assert message["type"] == "lifespan.shutdown" diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index 44a65c8b6..a9cb73e3a 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -166,7 +166,7 @@ async def asgi3app(scope, receive, send): assert scope == { "type": "lifespan", "asgi": {"version": "3.0", "spec_version": "2.0"}, - "extensions": {"state": {}}, + "state": {}, } async def test(): @@ -189,7 +189,7 @@ def asgi2app(scope): assert scope == { "type": "lifespan", "asgi": {"version": "2.0", "spec_version": "2.0"}, - "extensions": {"state": {}}, + "state": {}, } async def asgi(receive, send): @@ -254,7 +254,7 @@ async def app(scope, receive, send): message = await receive() assert message["type"] == "lifespan.startup" await send({"type": "lifespan.startup.complete"}) - scope["extensions"]["state"]["foo"] = 123 + scope["state"]["foo"] = 123 message = await receive() assert message["type"] == "lifespan.shutdown" await send({"type": "lifespan.shutdown.complete"}) diff --git a/uvicorn/lifespan/on.py b/uvicorn/lifespan/on.py index 7bcc3725a..37e935f01 100644 --- a/uvicorn/lifespan/on.py +++ b/uvicorn/lifespan/on.py @@ -83,7 +83,7 @@ async def main(self) -> None: scope: LifespanScope = { # type: ignore[typeddict-item] "type": "lifespan", "asgi": {"version": self.config.asgi_version, "spec_version": "2.0"}, - "extensions": {"state": self.state}, + "state": self.state, } await app(scope, self.receive, self.send) except BaseException as exc: From f6fe1aeb0ebd5c152ad44ab5ca380c3bd99b8d2c Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sat, 4 Mar 2023 12:48:58 -0600 Subject: [PATCH 19/22] Update tests/protocols/test_http.py Co-authored-by: Marcelo Trylesinski --- tests/protocols/test_http.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/protocols/test_http.py b/tests/protocols/test_http.py index dd31a5bb4..7ac1fb70c 100644 --- a/tests/protocols/test_http.py +++ b/tests/protocols/test_http.py @@ -999,8 +999,6 @@ async def app(scope, receive, send): @pytest.mark.anyio @pytest.mark.parametrize("protocol_cls", HTTP_PROTOCOLS) async def test_lifespan_state(protocol_cls): - app = Response("Hello, world", media_type="text/plain") - expected_states = [{"a": 123, "b": [1]}, {"a": 123, "b": [1, 2]}] async def app(scope, receive, send): From 86760bc526cb6708f9c7ea46bbd68b36be2c9318 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sat, 4 Mar 2023 13:02:11 -0600 Subject: [PATCH 20/22] move assertion out of loop --- tests/protocols/test_websocket.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index e41c44e6f..4578d3c6c 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -1,4 +1,5 @@ import asyncio +from copy import deepcopy import httpx import pytest @@ -1076,6 +1077,10 @@ async def test_lifespan_state(ws_protocol_cls, http_protocol_cls, unused_tcp_por {"a": 123, "b": [1, 2]}, ] + actual_states = [] + + state_equals_expected_state = True + async def lifespan_app(scope, receive, send): message = await receive() assert message["type"] == "lifespan.startup" @@ -1088,8 +1093,8 @@ async def lifespan_app(scope, receive, send): class App(WebSocketResponse): async def websocket_connect(self, message): - expected_state = expected_states.pop(0) - assert self.scope["state"] == expected_state + nonlocal state_equals_expected_state + actual_states.append(deepcopy(self.scope["state"])) self.scope["state"]["a"] = 456 self.scope["state"]["b"].append(2) await self.send({"type": "websocket.accept"}) @@ -1117,4 +1122,4 @@ async def app_wrapper(scope, receive, send): is_open = await open_connection(f"ws://127.0.0.1:{unused_tcp_port}") assert is_open - assert not expected_states # consumed + assert expected_states == actual_states From 79d31962b3f7bbbb92b2cccd4cc74b9eaa0df491 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sat, 4 Mar 2023 15:20:14 -0600 Subject: [PATCH 21/22] remove nonlocal --- tests/protocols/test_websocket.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index 6a5952b4f..7c6a427eb 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -1115,7 +1115,6 @@ async def lifespan_app(scope, receive, send): class App(WebSocketResponse): async def websocket_connect(self, message): - nonlocal state_equals_expected_state actual_states.append(deepcopy(self.scope["state"])) self.scope["state"]["a"] = 456 self.scope["state"]["b"].append(2) From 40dcfabb716c61e8ee3cfb9bbc22661cf6e37202 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Sun, 5 Mar 2023 14:30:12 +0100 Subject: [PATCH 22/22] Update tests/protocols/test_websocket.py --- tests/protocols/test_websocket.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index 7c6a427eb..6bc56dfa4 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -1101,8 +1101,6 @@ async def test_lifespan_state(ws_protocol_cls, http_protocol_cls, unused_tcp_por actual_states = [] - state_equals_expected_state = True - async def lifespan_app(scope, receive, send): message = await receive() assert message["type"] == "lifespan.startup"