1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
|
from __future__ import annotations
import re
from collections import defaultdict
from functools import lru_cache
from traceback import format_exc
from typing import TYPE_CHECKING, Any, Pattern
from litestar._asgi.routing_trie import validate_node
from litestar._asgi.routing_trie.mapping import add_route_to_trie
from litestar._asgi.routing_trie.traversal import parse_path_to_route
from litestar._asgi.routing_trie.types import create_node
from litestar._asgi.utils import get_route_handlers
from litestar.exceptions import ImproperlyConfiguredException
from litestar.utils import normalize_path
__all__ = ("ASGIRouter",)
if TYPE_CHECKING:
from litestar._asgi.routing_trie.types import RouteTrieNode
from litestar.app import Litestar
from litestar.routes import ASGIRoute, HTTPRoute, WebSocketRoute
from litestar.routes.base import BaseRoute
from litestar.types import (
ASGIApp,
LifeSpanReceive,
LifeSpanSend,
LifeSpanShutdownCompleteEvent,
LifeSpanShutdownFailedEvent,
LifeSpanStartupCompleteEvent,
LifeSpanStartupFailedEvent,
Method,
Receive,
RouteHandlerType,
Scope,
Send,
)
class ASGIRouter:
"""Litestar ASGI router.
Handling both the ASGI lifespan events and routing of connection requests.
"""
__slots__ = (
"_mount_paths_regex",
"_mount_routes",
"_plain_routes",
"_registered_routes",
"_static_routes",
"app",
"root_route_map_node",
"route_handler_index",
"route_mapping",
)
def __init__(self, app: Litestar) -> None:
"""Initialize ``ASGIRouter``.
Args:
app: The Litestar app instance
"""
self._mount_paths_regex: Pattern | None = None
self._mount_routes: dict[str, RouteTrieNode] = {}
self._plain_routes: set[str] = set()
self._registered_routes: set[HTTPRoute | WebSocketRoute | ASGIRoute] = set()
self.app = app
self.root_route_map_node: RouteTrieNode = create_node()
self.route_handler_index: dict[str, RouteHandlerType] = {}
self.route_mapping: dict[str, list[BaseRoute]] = defaultdict(list)
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
"""ASGI callable.
The main entry point to the Router class.
"""
scope.setdefault("path_params", {})
path = scope["path"]
if root_path := scope.get("root_path", ""):
path = path.split(root_path, maxsplit=1)[-1]
normalized_path = normalize_path(path)
asgi_app, scope["route_handler"], scope["path"], scope["path_params"] = self.handle_routing(
path=normalized_path, method=scope.get("method")
)
await asgi_app(scope, receive, send)
@lru_cache(1024) # noqa: B019
def handle_routing(self, path: str, method: Method | None) -> tuple[ASGIApp, RouteHandlerType, str, dict[str, Any]]:
"""Handle routing for a given path / method combo. This method is meant to allow easy caching.
Args:
path: The path of the request.
method: The scope's method, if any.
Returns:
A tuple composed of the ASGIApp of the route, the route handler instance, the resolved and normalized path and any parsed path params.
"""
return parse_path_to_route(
mount_paths_regex=self._mount_paths_regex,
mount_routes=self._mount_routes,
path=path,
plain_routes=self._plain_routes,
root_node=self.root_route_map_node,
method=method,
)
def _store_handler_to_route_mapping(self, route: BaseRoute) -> None:
"""Store the mapping of route handlers to routes and to route handler names.
Args:
route: A Route instance.
Returns:
None
"""
for handler in get_route_handlers(route):
if handler.name in self.route_handler_index and str(self.route_handler_index[handler.name]) != str(handler):
raise ImproperlyConfiguredException(
f"route handler names must be unique - {handler.name} is not unique."
)
identifier = handler.name or str(handler)
self.route_mapping[identifier].append(route)
self.route_handler_index[identifier] = handler
def construct_routing_trie(self) -> None:
"""Create a map of the app's routes.
This map is used in the asgi router to route requests.
"""
new_routes = [route for route in self.app.routes if route not in self._registered_routes]
for route in new_routes:
add_route_to_trie(
app=self.app,
mount_routes=self._mount_routes,
plain_routes=self._plain_routes,
root_node=self.root_route_map_node,
route=route,
)
self._store_handler_to_route_mapping(route)
self._registered_routes.add(route)
validate_node(node=self.root_route_map_node)
if self._mount_routes:
self._mount_paths_regex = re.compile("|".join(sorted(set(self._mount_routes)))) # pyright: ignore
async def lifespan(self, receive: LifeSpanReceive, send: LifeSpanSend) -> None:
"""Handle the ASGI "lifespan" event on application startup and shutdown.
Args:
receive: The ASGI receive function.
send: The ASGI send function.
Returns:
None.
"""
message = await receive()
shutdown_event: LifeSpanShutdownCompleteEvent = {"type": "lifespan.shutdown.complete"}
startup_event: LifeSpanStartupCompleteEvent = {"type": "lifespan.startup.complete"}
try:
async with self.app.lifespan():
await send(startup_event)
message = await receive()
except BaseException as e:
formatted_exception = format_exc()
failure_message: LifeSpanStartupFailedEvent | LifeSpanShutdownFailedEvent
if message["type"] == "lifespan.startup":
failure_message = {"type": "lifespan.startup.failed", "message": formatted_exception}
else:
failure_message = {"type": "lifespan.shutdown.failed", "message": formatted_exception}
await send(failure_message)
raise e
await send(shutdown_event)
|
