26 lines
928 B
Python
26 lines
928 B
Python
from time import perf_counter
|
|
|
|
from starlette.datastructures import MutableHeaders
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
|
|
class ResponseTimeMiddleware:
|
|
def __init__(self, app: ASGIApp, header_name: str = "Elapsed") -> None:
|
|
self.app = app
|
|
self.header_name = header_name
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
start = perf_counter()
|
|
|
|
if scope["type"] not in ("http", "websocket"):
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
async def send_wrapper(message: Message) -> None:
|
|
if message["type"] == "http.response.start":
|
|
headers = MutableHeaders(scope=message)
|
|
elapsed = perf_counter() - start
|
|
headers.append(self.header_name, str(elapsed))
|
|
await send(message)
|
|
|
|
await self.app(scope, receive, send_wrapper)
|