Skip to content
Open
2 changes: 2 additions & 0 deletions CHANGES/11483.feature.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Added ``StreamReader.total_raw_bytes`` to check the number of bytes downloaded
-- by :user:`robpats`.
1 change: 1 addition & 0 deletions CONTRIBUTORS.txt
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,7 @@ Pahaz Blinov
Panagiotis Kolokotronis
Pankaj Pandey
Parag Jain
Patrick Lee
Pau Freixes
Paul Colomiets
Paul J. Dorn
Expand Down
2 changes: 2 additions & 0 deletions aiohttp/http_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -936,6 +936,7 @@ class DeflateBuffer:
def __init__(self, out: StreamReader, encoding: Optional[str]) -> None:
self.out = out
self.size = 0
self.out.total_compressed_bytes = self.size
self.encoding = encoding
self._started_decoding = False

Expand Down Expand Up @@ -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"
Expand Down
8 changes: 8 additions & 0 deletions aiohttp/streams.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,7 @@ class StreamReader(AsyncStreamReaderMixin):
"_eof_callbacks",
"_eof_counter",
"total_bytes",
"total_compressed_bytes",
)

def __init__(
Expand Down Expand Up @@ -159,6 +160,7 @@ def __init__(
self._eof_callbacks: List[Callable[[], None]] = []
self._eof_counter = 0
self.total_bytes = 0
self.total_compressed_bytes = None

def __repr__(self) -> str:
info = [self.__class__.__name__]
Expand Down Expand Up @@ -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(
Expand Down
9 changes: 8 additions & 1 deletion docs/streams.rst
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ Streaming API
:attr:`aiohttp.ClientResponse.content` properties for accessing raw
BODY data.

Reading Methods
Reading Attributes and Methods
---------------

.. method:: StreamReader.read(n=-1)
Expand Down Expand Up @@ -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
------------------------------

Expand Down
43 changes: 43 additions & 0 deletions tests/test_client_functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
Loading