Add output_size property (#12452)

Author: Dreamsorcerer

CHANGES/12452.feature.rst

@@ -0,0 +1,2 @@
+Added :attr:`~aiohttp.ClientResponse.output_size` and
+:attr:`~aiohttp.ClientResponse.upload_complete` -- by :user:`Dreamsorcerer`.

aiohttp/client_reqrep.py

@@ -212,6 +212,9 @@ class ClientResponse(HeadersMixin):
     _resolve_charset: Callable[["ClientResponse", bytes], str] = lambda *_: "utf-8"
 
     __writer: asyncio.Task[None] | None = None
+    _stream_writer: AbstractStreamWriter | None = None
+    _output_size: int = 0
+    _upload_complete: asyncio.Future[None] | None = None
 
     def __init__(
         self,
@@ -226,6 +229,7 @@ def __init__(
         session: "ClientSession | None",
         request_headers: CIMultiDict[str],
         original_url: URL,
+        stream_writer: AbstractStreamWriter,
         **kwargs: object,
     ) -> None:
         # kwargs exists so authors of subclasses should expect to pass through unknown
@@ -240,7 +244,10 @@ def __init__(
 
         self._real_url = url
         self._url = url.with_fragment(None) if url.raw_fragment else url
-        if writer is not None:
+        if writer is None:  # Request already sent
+            self._output_size = stream_writer.output_size
+        else:
+            self._stream_writer = stream_writer
             self._writer = writer
         if continue100 is not None:
             self._continue = continue100
@@ -261,6 +268,11 @@ def __init__(
 
     def __reset_writer(self, _: object = None) -> None:
         self.__writer = None
+        if self._stream_writer is not None:
+            self._output_size = self._stream_writer.output_size
+            self._stream_writer = None
+        if self._upload_complete is not None and not self._upload_complete.done():
+            self._upload_complete.set_result(None)
 
     @property
     def _writer(self) -> asyncio.Task[None] | None:
@@ -281,10 +293,29 @@ def _writer(self, writer: asyncio.Task[None] | None) -> None:
             return
         if writer.done():
             # The writer is already done, so we can clear it immediately.
-            self.__writer = None
+            self.__reset_writer()
         else:
             writer.add_done_callback(self.__reset_writer)
 
+    @property
+    def output_size(self) -> int:
+        """Number of bytes sent for this request."""
+        if self._stream_writer is not None:
+            return self._stream_writer.output_size
+        return self._output_size
+
+    @property
+    def upload_complete(self) -> "asyncio.Future[None]":
+        """Future set when the request body has been fully sent.
+
+        Already done when the request had no body or was written eagerly.
+        """
+        if self._upload_complete is None:
+            self._upload_complete = self._loop.create_future()
+            if self._stream_writer is None:  # upload already finished
+                self._upload_complete.set_result(None)
+        return self._upload_complete
+
     @property
     def cookies(self) -> SimpleCookie:
         if self._cookies is None:
@@ -558,6 +589,9 @@ async def _wait_released(self) -> None:
     def _cleanup_writer(self) -> None:
         if self.__writer is not None:
             self.__writer.cancel()
+        if self._stream_writer is not None:
+            self._output_size = self._stream_writer.output_size
+            self._stream_writer = None
         self._session = None
 
     def _notify_content(self) -> None:
@@ -800,7 +834,11 @@ def _update_headers(self, headers: CIMultiDict[str]) -> None:
         self.headers[hdrs.HOST] = headers.pop(hdrs.HOST, host)
         self.headers.extend(headers)
 
-    def _create_response(self, task: asyncio.Task[None] | None) -> ClientResponse:
+    def _create_response(
+        self,
+        task: asyncio.Task[None] | None,
+        stream_writer: AbstractStreamWriter,
+    ) -> ClientResponse:
         return self.response_class(
             self.method,
             self.original_url,
@@ -812,6 +850,7 @@ def _create_response(self, task: asyncio.Task[None] | None) -> ClientResponse:
             session=None,
             request_headers=self.headers,
             original_url=self.original_url,
+            stream_writer=stream_writer,
         )
 
     def _create_writer(self, protocol: BaseProtocol) -> StreamWriter:
@@ -885,7 +924,7 @@ async def _send(self, conn: "Connection") -> ClientResponse:
             protocol.start_timeout()
             writer.set_eof()
             task = None
-        self._response = self._create_response(task)
+        self._response = self._create_response(task, stream_writer=writer)
         return self._response
 
     async def _write_bytes(
@@ -1261,7 +1300,11 @@ def _update_proxy(
         self.proxy = proxy
         self.proxy_headers = proxy_headers
 
-    def _create_response(self, task: asyncio.Task[None] | None) -> ClientResponse:
+    def _create_response(
+        self,
+        task: asyncio.Task[None] | None,
+        stream_writer: AbstractStreamWriter,
+    ) -> ClientResponse:
         return self.response_class(
             self.method,
             self.original_url,
@@ -1273,6 +1316,7 @@ def _create_response(self, task: asyncio.Task[None] | None) -> ClientResponse:
             session=self._session,
             request_headers=self.headers,
             original_url=self.original_url,
+            stream_writer=stream_writer,
         )
 
     def _create_writer(self, protocol: BaseProtocol) -> StreamWriter:

docs/client_reference.rst

@@ -1538,6 +1538,30 @@ Response object
 
       .. versionadded:: 3.2
 
+   .. attribute:: output_size
+
+      Number of bytes sent for this request.
+
+      Pair with :attr:`upload_complete` to display upload progress::
+
+          async with session.post(url, data=mpwriter) as resp:
+              while not resp.upload_complete.done():
+                  print(f"uploaded {resp.output_size} bytes")
+                  await asyncio.sleep(0.5)
+              print(f"upload complete: {resp.output_size} bytes")
+
+      .. versionadded:: 3.14
+
+   .. attribute:: upload_complete
+
+      An :class:`asyncio.Future` set when the request body has been fully sent.
+
+      Use ``await resp.upload_complete`` to block until the upload finishes, or
+      ``resp.upload_complete.done()`` to poll from a progress-sampling loop
+      (see :attr:`output_size`).
+
+      .. versionadded:: 3.14
+
    .. attribute:: content_type
 
       Read-only property with *content* part of *Content-Type* header.

tests/test_client_functional.py

@@ -5746,3 +5746,173 @@ async def handler(request: web.Request) -> web.Response:
         data = await resp.content.read()
         assert resp.content.total_raw_bytes == len(data)
         assert resp.content.total_raw_bytes == int(resp.headers["Content-Length"])
+
+
+async def test_output_size_bytes(aiohttp_client: AiohttpClient) -> None:
+    async def handler(request: web.Request) -> web.Response:
+        await request.read()
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_post("/", handler)
+    client = await aiohttp_client(app)
+
+    body = b"x" * 1024
+    async with client.post("/", data=body) as resp:
+        assert resp.output_size >= len(body)
+
+
+async def test_output_size_multipart(aiohttp_client: AiohttpClient) -> None:
+    async def handler(request: web.Request) -> web.Response:
+        await request.read()
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_post("/", handler)
+    client = await aiohttp_client(app)
+
+    mpwriter = aiohttp.MultipartWriter("form-data")
+    mpwriter.append(b"x" * 4096)
+    mpwriter.append(b"y" * 2048)
+    expected_body_size = mpwriter.size
+    assert expected_body_size is not None
+
+    async with client.post("/", data=mpwriter) as resp:
+        assert resp.output_size >= expected_body_size
+
+
+async def test_output_size_keepalive_isolated(
+    aiohttp_client: AiohttpClient,
+) -> None:
+    """Each request on a keep-alive connection has its own counter."""
+    transports: set[object] = set()
+
+    async def handler(request: web.Request) -> web.Response:
+        transports.add(request.transport)
+        await request.read()
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_post("/", handler)
+    connector = aiohttp.TCPConnector(limit=1, force_close=False)
+    client = await aiohttp_client(app, connector=connector)
+    body = b"x" * 65536
+
+    async with client.post("/", data=body) as resp1:
+        size1 = resp1.output_size
+
+    async with client.post("/", data=body) as resp2:
+        size2 = resp2.output_size
+
+    assert len(transports) == 1  # Check keep-alive worked.
+    assert size1 >= len(body)
+    assert size1 == size2
+
+
+async def test_output_size_progress(aiohttp_client: AiohttpClient) -> None:
+    """output_size advances by exactly one chunk per yield."""
+
+    async def handler(request: web.Request) -> web.StreamResponse:
+        response = web.StreamResponse()
+        await response.prepare(request)
+        # Flush headers + a chunk so resp.start() returns on the client
+        # side before we read the body.
+        await response.write(b"x")
+        await request.read()
+        return response
+
+    app = web.Application()
+    app.router.add_post("/", handler)
+    client = await aiohttp_client(app)
+
+    chunk_size = 4096
+    chunk = b"z" * chunk_size
+    num_chunks = 8
+    sample_taken = asyncio.Event()
+    next_chunk = asyncio.Event()
+
+    async def gated_body() -> AsyncIterator[bytes]:
+        for _ in range(num_chunks):
+            yield chunk
+            sample_taken.clear()
+            next_chunk.set()
+            await sample_taken.wait()
+
+    async with client.post("/", data=gated_body()) as resp:
+        samples: list[int] = []
+        for _ in range(num_chunks):
+            await next_chunk.wait()
+            next_chunk.clear()
+            samples.append(resp.output_size)
+            assert not resp.upload_complete.done()
+            sample_taken.set()
+        await resp.upload_complete
+        assert resp.upload_complete.done()
+        await resp.read()
+
+    # Each sample after the first reflects exactly one more chunk on the wire.
+    chunked_framing = len(f"{chunk_size:x}".encode()) + 4
+    deltas = [samples[i] - samples[i - 1] for i in range(1, len(samples))]
+    assert deltas == [chunk_size + chunked_framing] * (num_chunks - 1)
+
+
+async def test_output_size_get_request(aiohttp_client: AiohttpClient) -> None:
+    """GET request with no body still reports the request header byte count."""
+
+    async def handler(request: web.Request) -> web.Response:
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_get("/", handler)
+    client = await aiohttp_client(app)
+
+    async with client.get("/") as resp:
+        assert resp.output_size >= 0
+
+
+async def test_output_size_writer_released(aiohttp_client: AiohttpClient) -> None:
+    """Writer is dropped once body upload completes; output_size survives."""
+
+    async def handler(request: web.Request) -> web.Response:
+        await request.read()
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_post("/", handler)
+    client = await aiohttp_client(app)
+
+    body = b"x" * 1024
+    async with client.post("/", data=body) as resp:
+        await resp.read()
+        assert resp._stream_writer is None
+    assert resp.output_size >= len(body)
+
+
+async def test_upload_complete_no_body(aiohttp_client: AiohttpClient) -> None:
+    async def handler(request: web.Request) -> web.Response:
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_get("/", handler)
+    client = await aiohttp_client(app)
+
+    async with client.get("/") as resp:
+        assert resp.upload_complete.done()
+
+
+async def test_upload_complete_late_access(aiohttp_client: AiohttpClient) -> None:
+    """Accessing upload_complete after the upload finished returns a done future."""
+
+    async def handler(request: web.Request) -> web.Response:
+        await request.read()
+        return web.Response()
+
+    app = web.Application()
+    app.router.add_post("/", handler)
+    client = await aiohttp_client(app)
+
+    async with client.post("/", data=b"x" * 1024) as resp:
+        await resp.read()
+        # Writer task is done; future is created lazily on this first access.
+        assert resp._upload_complete is None
+        assert resp.upload_complete.done()

tests/test_client_proto.py

@@ -6,6 +6,7 @@
 from yarl import URL
 
 from aiohttp import http
+from aiohttp.abc import AbstractStreamWriter
 from aiohttp.client_exceptions import ClientOSError, ServerDisconnectedError
 from aiohttp.client_proto import ResponseHandler
 from aiohttp.client_reqrep import ClientResponse
@@ -122,6 +123,9 @@ async def test_multiple_responses_one_byte_at_a_time() -> None:
                 session=mock.Mock(),
                 request_headers=CIMultiDict[str](),
                 original_url=url,
+                stream_writer=mock.create_autospec(
+                    AbstractStreamWriter, spec_set=True, instance=True
+                ),
             )
             await response.start(conn)
             await response.read() == payload
@@ -153,6 +157,9 @@ class PatchableHttpResponseParser(http.HttpResponseParser):
             session=mock.Mock(),
             request_headers=CIMultiDict[str](),
             original_url=url,
+            stream_writer=mock.create_autospec(
+                AbstractStreamWriter, spec_set=True, instance=True
+            ),
         )
         await response.start(conn)
         await response.read() == b"ab"
@@ -184,6 +191,9 @@ async def test_client_protocol_readuntil_eof() -> None:
         session=mock.Mock(),
         request_headers=CIMultiDict[str](),
         original_url=url,
+        stream_writer=mock.create_autospec(
+            AbstractStreamWriter, spec_set=True, instance=True
+        ),
     )
     proto.set_response_params(read_until_eof=True)
     await response.start(conn)