multipart: aclosing support

Author: tsimoshka

httpx/_compat.py

@@ -60,4 +60,23 @@ def set_minimum_tls_version_1_2(context: ssl.SSLContext) -> None:
         context.options |= ssl.OP_NO_TLSv1_1
 
 
-__all__ = ["brotli", "set_minimum_tls_version_1_2"]
+if sys.version_info >= (3, 10):
+    from contextlib import aclosing
+else:
+    from contextlib import asynccontextmanager
+    from typing import Awaitable, Protocol, TypeVar
+
+    class _SupportsAclose(Protocol):
+        def aclose(self) -> Awaitable[object]: ...
+
+    _SupportsAcloseT = TypeVar("_SupportsAcloseT", bound=_SupportsAclose)
+
+    @asynccontextmanager
+    async def aclosing(thing: _SupportsAcloseT) -> None:
+        try:
+            yield thing
+        finally:
+            await thing.aclose()
+
+
+__all__ = ["brotli", "set_minimum_tls_version_1_2", "aclosing"]

httpx/_multipart.py

@@ -6,6 +6,7 @@
 from collections.abc import AsyncIterable
 from pathlib import Path
 
+from ._compat import aclosing
 from ._types import (
     AsyncByteStream,
     FileContent,
@@ -188,7 +189,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
@@ -216,10 +217,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):
@@ -263,12 +265,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
@@ -309,5 +312,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