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
|
from __future__ import annotations
from asyncio import sleep
from litestar import Litestar, get
from litestar.middleware import MiddlewareProtocol
from litestar.types import ASGIApp, Message, Receive, Scope, Send
from pyinstrument import Profiler
class ProfilingMiddleware(MiddlewareProtocol):
def __init__(self, app: ASGIApp) -> None:
super().__init__(app) # type: ignore
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
profiler = Profiler(interval=0.001, async_mode="enabled")
profiler.start()
profile_html: str | None = None
async def send_wrapper(message: Message) -> None:
if message["type"] == "http.response.start":
profiler.stop()
nonlocal profile_html
profile_html = profiler.output_html()
message["headers"] = [
(b"content-type", b"text/html; charset=utf-8"),
(b"content-length", str(len(profile_html)).encode()),
]
elif message["type"] == "http.response.body":
assert profile_html is not None
message["body"] = profile_html.encode()
await send(message)
await self.app(scope, receive, send_wrapper)
@get("/")
async def index() -> str:
await sleep(1)
return "Hello, world!"
app = Litestar(
route_handlers=[index],
middleware=[ProfilingMiddleware],
)
|
