fix ensure_next read first chunk

abersheeran · Jul 26, 2024 · 3635cc3 · 3635cc3
1 parent 27f93ab
commit 3635cc3
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/baize/wsgi/middleware.py b/baize/wsgi/middleware.py
@@ -19,8 +19,13 @@ def stream(self, chunk_size: int = 4096 * 16) -> Iterator[bytes]:
 
 
 def ensure_next(iterable: Iterable[bytes]) -> Iterable[bytes]:
-    yield iterable.__iter__().__next__()
-    yield from iterable
+    first_chunk = iterable.__iter__().__next__()
+
+    def generator():
+        yield first_chunk
+        yield from iterable
+
+    return generator()
 
 
 class NextResponse(StreamingResponse):