multipart: aclosing support

encode · Dec 1, 2024 · 7984e22 · 7984e22
1 parent 15de00c
commit 7984e22
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/httpx/_compat.py b/httpx/_compat.py
diff --git a/httpx/_multipart.py b/httpx/_multipart.py
@@ -8,6 +8,7 @@
 from collections.abc import AsyncIterable
 from pathlib import Path
 
+from ._compat import aclosing
 from ._types import (
     AsyncByteStream,
     FileContent,
@@ -219,7 +220,7 @@ def render_data(self) -> typing.Iterator[bytes]:
             yield to_bytes(chunk)
             chunk = self.file.read(self.CHUNK_SIZE)
 
-    async def arender_data(self) -> typing.AsyncIterator[bytes]:
+    async def arender_data(self) -> typing.AsyncGenerator[bytes]:
         if not isinstance(self.file, AsyncIterable):
             for chunk in self.render_data():
                 yield chunk
@@ -247,10 +248,11 @@ def render(self) -> typing.Iterator[bytes]:
         yield self.render_headers()
         yield from self.render_data()
 
-    async def arender(self) -> typing.AsyncIterator[bytes]:
+    async def arender(self) -> typing.AsyncGenerator[bytes]:
         yield self.render_headers()
-        async for chunk in self.arender_data():
-            yield chunk
+        async with aclosing(self.arender_data()) as data:
+            async for chunk in data:
+                yield chunk
 
 
 class MultipartStream(SyncByteStream, AsyncByteStream):
@@ -294,12 +296,13 @@ def iter_chunks(self) -> typing.Iterator[bytes]:
             yield b"\r\n"
         yield b"--%s--\r\n" % self.boundary
 
-    async def aiter_chunks(self) -> typing.AsyncIterator[bytes]:
+    async def aiter_chunks(self) -> typing.AsyncGenerator[bytes]:
         for field in self.fields:
             yield b"--%s\r\n" % self.boundary
             if isinstance(field, FileField):
-                async for chunk in field.arender():
-                    yield chunk
+                async with aclosing(field.arender()) as data:
+                    async for chunk in data:
+                        yield chunk
             else:
                 for chunk in field.render():
                     yield chunk
@@ -340,5 +343,6 @@ def __iter__(self) -> typing.Iterator[bytes]:
             yield chunk
 
     async def __aiter__(self) -> typing.AsyncIterator[bytes]:
-        async for chunk in self.aiter_chunks():
-            yield chunk
+        async with aclosing(self.aiter_chunks()) as data:
+            async for chunk in data:
+                yield chunk