Add StreamReader.total_raw_bytes to report the download progress (#11483)

Author: robpats

CHANGES/11483.feature.rst

@@ -0,0 +1,2 @@
+Added ``StreamReader.total_raw_bytes`` to check the number of bytes downloaded
+-- by :user:`robpats`.

CONTRIBUTORS.txt

@@ -287,6 +287,7 @@ Pahaz Blinov
 Panagiotis Kolokotronis
 Pankaj Pandey
 Parag Jain
+Patrick Lee
 Pau Freixes
 Paul Colomiets
 Paul J. Dorn

aiohttp/http_parser.py

@@ -936,6 +936,7 @@ class DeflateBuffer:
     def __init__(self, out: StreamReader, encoding: Optional[str]) -> None:
         self.out = out
         self.size = 0
+        out.total_compressed_bytes = self.size
         self.encoding = encoding
         self._started_decoding = False
 
@@ -969,6 +970,7 @@ def feed_data(self, chunk: bytes) -> None:
             return
 
         self.size += len(chunk)
+        self.out.total_compressed_bytes = self.size
 
         # RFC1950
         # bits 0..3 = CM = 0b1000 = 8 = "deflate"

aiohttp/streams.py

@@ -132,6 +132,7 @@ class StreamReader(AsyncStreamReaderMixin):
         "_eof_callbacks",
         "_eof_counter",
         "total_bytes",
+        "total_compressed_bytes",
     )
 
     def __init__(
@@ -159,6 +160,7 @@ def __init__(
         self._eof_callbacks: List[Callable[[], None]] = []
         self._eof_counter = 0
         self.total_bytes = 0
+        self.total_compressed_bytes: Optional[int] = None
 
     def __repr__(self) -> str:
         info = [self.__class__.__name__]
@@ -250,6 +252,12 @@ async def wait_eof(self) -> None:
         finally:
             self._eof_waiter = None
 
+    @property
+    def total_raw_bytes(self) -> int:
+        if self.total_compressed_bytes is None:
+            return self.total_bytes
+        return self.total_compressed_bytes
+
     def unread_data(self, data: bytes) -> None:
         """rollback reading some data from stream, inserting it to buffer head."""
         warnings.warn(

docs/streams.rst

@@ -20,8 +20,8 @@ Streaming API
    :attr:`aiohttp.ClientResponse.content` properties for accessing raw
    BODY data.
 
-Reading Methods
----------------
+Reading Attributes and Methods
+------------------------------
 
 .. method:: StreamReader.read(n=-1)
       :async:
@@ -109,6 +109,13 @@ Reading Methods
                                to the end of a HTTP chunk.
 
 
+.. attribute:: StreamReader.total_raw_bytes
+
+   The number of bytes of raw data downloaded (before decompression).
+
+   Readonly :class:`int` property.
+
+
 Asynchronous Iteration Support
 ------------------------------
 

tests/test_client_functional.py

@@ -5586,3 +5586,46 @@ async def handler(request: web.Request) -> web.Response:
 
     finally:
         await asyncio.to_thread(f.close)
+
+
+async def test_stream_reader_total_raw_bytes(aiohttp_client: AiohttpClient) -> None:
+    """Test whether StreamReader.total_raw_bytes returns the number of bytes downloaded"""
+    source_data = b"@dKal^pH>1h|YW1:c2J$" * 4096
+
+    async def handler(request: web.Request) -> web.Response:
+        response = web.Response(body=source_data)
+        response.enable_compression()
+        return response
+
+    app = web.Application()
+    app.router.add_get("/", handler)
+
+    client = await aiohttp_client(app)
+
+    # Check for decompressed data
+    async with client.get(
+        "/", headers={"Accept-Encoding": "gzip"}, auto_decompress=True
+    ) as resp:
+        assert resp.headers["Content-Encoding"] == "gzip"
+        assert int(resp.headers["Content-Length"]) < len(source_data)
+        data = await resp.content.read()
+        assert len(data) == len(source_data)
+        assert resp.content.total_raw_bytes == int(resp.headers["Content-Length"])
+
+    # Check for compressed data
+    async with client.get(
+        "/", headers={"Accept-Encoding": "gzip"}, auto_decompress=False
+    ) as resp:
+        assert resp.headers["Content-Encoding"] == "gzip"
+        data = await resp.content.read()
+        assert resp.content.total_raw_bytes == len(data)
+        assert resp.content.total_raw_bytes == int(resp.headers["Content-Length"])
+
+    # Check for non-compressed data
+    async with client.get(
+        "/", headers={"Accept-Encoding": "identity"}, auto_decompress=True
+    ) as resp:
+        assert "Content-Encoding" not in resp.headers
+        data = await resp.content.read()
+        assert resp.content.total_raw_bytes == len(data)
+        assert resp.content.total_raw_bytes == int(resp.headers["Content-Length"])