Skip to content

Commit ce38e43

Browse files
Add StreamReader.total_raw_bytes to report the download progress (#11483) (#11576)
(cherry picked from commit 2c9dbad) Co-authored-by: robpats <[email protected]>
1 parent f3cd943 commit ce38e43

File tree

6 files changed

+65
-2
lines changed

6 files changed

+65
-2
lines changed

CHANGES/11483.feature.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Added ``StreamReader.total_raw_bytes`` to check the number of bytes downloaded
2+
-- by :user:`robpats`.

CONTRIBUTORS.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,7 @@ Pahaz Blinov
279279
Panagiotis Kolokotronis
280280
Pankaj Pandey
281281
Parag Jain
282+
Patrick Lee
282283
Pau Freixes
283284
Paul Colomiets
284285
Paul J. Dorn

aiohttp/http_parser.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -964,6 +964,7 @@ class DeflateBuffer:
964964
def __init__(self, out: StreamReader, encoding: Optional[str]) -> None:
965965
self.out = out
966966
self.size = 0
967+
out.total_compressed_bytes = self.size
967968
self.encoding = encoding
968969
self._started_decoding = False
969970

@@ -997,6 +998,7 @@ def feed_data(self, chunk: bytes, size: int) -> None:
997998
return
998999

9991000
self.size += size
1001+
self.out.total_compressed_bytes = self.size
10001002

10011003
# RFC1950
10021004
# bits 0..3 = CM = 0b1000 = 8 = "deflate"

aiohttp/streams.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ class StreamReader(AsyncStreamReaderMixin):
130130
"_eof_callbacks",
131131
"_eof_counter",
132132
"total_bytes",
133+
"total_compressed_bytes",
133134
)
134135

135136
def __init__(
@@ -159,6 +160,7 @@ def __init__(
159160
self._eof_callbacks: List[Callable[[], None]] = []
160161
self._eof_counter = 0
161162
self.total_bytes = 0
163+
self.total_compressed_bytes: Optional[int] = None
162164

163165
def __repr__(self) -> str:
164166
info = [self.__class__.__name__]
@@ -250,6 +252,12 @@ async def wait_eof(self) -> None:
250252
finally:
251253
self._eof_waiter = None
252254

255+
@property
256+
def total_raw_bytes(self) -> int:
257+
if self.total_compressed_bytes is None:
258+
return self.total_bytes
259+
return self.total_compressed_bytes
260+
253261
def unread_data(self, data: bytes) -> None:
254262
"""rollback reading some data from stream, inserting it to buffer head."""
255263
warnings.warn(

docs/streams.rst

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ Streaming API
2020
:attr:`aiohttp.ClientResponse.content` properties for accessing raw
2121
BODY data.
2222

23-
Reading Methods
24-
---------------
23+
Reading Attributes and Methods
24+
------------------------------
2525

2626
.. method:: StreamReader.read(n=-1)
2727
:async:
@@ -109,6 +109,13 @@ Reading Methods
109109
to the end of a HTTP chunk.
110110

111111

112+
.. attribute:: StreamReader.total_raw_bytes
113+
114+
The number of bytes of raw data downloaded (before decompression).
115+
116+
Readonly :class:`int` property.
117+
118+
112119
Asynchronous Iteration Support
113120
------------------------------
114121

tests/test_client_functional.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5515,3 +5515,46 @@ async def handler(request: web.Request) -> web.Response:
55155515

55165516
finally:
55175517
await asyncio.to_thread(f.close)
5518+
5519+
5520+
async def test_stream_reader_total_raw_bytes(aiohttp_client: AiohttpClient) -> None:
5521+
"""Test whether StreamReader.total_raw_bytes returns the number of bytes downloaded"""
5522+
source_data = b"@dKal^pH>1h|YW1:c2J$" * 4096
5523+
5524+
async def handler(request: web.Request) -> web.Response:
5525+
response = web.Response(body=source_data)
5526+
response.enable_compression()
5527+
return response
5528+
5529+
app = web.Application()
5530+
app.router.add_get("/", handler)
5531+
5532+
client = await aiohttp_client(app)
5533+
5534+
# Check for decompressed data
5535+
async with client.get(
5536+
"/", headers={"Accept-Encoding": "gzip"}, auto_decompress=True
5537+
) as resp:
5538+
assert resp.headers["Content-Encoding"] == "gzip"
5539+
assert int(resp.headers["Content-Length"]) < len(source_data)
5540+
data = await resp.content.read()
5541+
assert len(data) == len(source_data)
5542+
assert resp.content.total_raw_bytes == int(resp.headers["Content-Length"])
5543+
5544+
# Check for compressed data
5545+
async with client.get(
5546+
"/", headers={"Accept-Encoding": "gzip"}, auto_decompress=False
5547+
) as resp:
5548+
assert resp.headers["Content-Encoding"] == "gzip"
5549+
data = await resp.content.read()
5550+
assert resp.content.total_raw_bytes == len(data)
5551+
assert resp.content.total_raw_bytes == int(resp.headers["Content-Length"])
5552+
5553+
# Check for non-compressed data
5554+
async with client.get(
5555+
"/", headers={"Accept-Encoding": "identity"}, auto_decompress=True
5556+
) as resp:
5557+
assert "Content-Encoding" not in resp.headers
5558+
data = await resp.content.read()
5559+
assert resp.content.total_raw_bytes == len(data)
5560+
assert resp.content.total_raw_bytes == int(resp.headers["Content-Length"])

0 commit comments

Comments
 (0)