fixup: pr

joerick · Jan 26, 2024 · 6ff912e · 6ff912e
1 parent 8ee43c1
commit 6ff912e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 14 deletions.
diff --git a/docs/guide.md b/docs/guide.md
@@ -284,20 +284,20 @@ class ProfilingMiddleware(MiddlewareProtocol):
 
     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.body":
-                message["body"] = profiler.output_html().encode()
-
-            elif message["type"] == "http.response.start":
+            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(profiler.output_html())).encode()),
+                    (b"content-length", str(len(profile_html)).encode()),
                 ]
-
+            elif message["type"] == "http.response.body":
+                message["body"] = profile_html.encode()
             await send(message)
 
         await self.app(scope, receive, send_wrapper)

diff --git a/examples/litestar_hello.py b/examples/litestar_hello.py
@@ -13,20 +13,20 @@ def __init__(self, app: ASGIApp) -> None:
 
     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.body":
-                message["body"] = profiler.output_html().encode()
-
-            elif message["type"] == "http.response.start":
+            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(profiler.output_html())).encode()),
+                    (b"content-length", str(len(profile_html)).encode()),
                 ]
-
+            elif message["type"] == "http.response.body":
+                message["body"] = profile_html.encode()
             await send(message)
 
         await self.app(scope, receive, send_wrapper)