Skip to content

Commit

Permalink
Read stream for downloads
Browse files Browse the repository at this point in the history
  • Loading branch information
BdEgh committed Mar 17, 2024
1 parent 976f04f commit e54eb53
Show file tree
Hide file tree
Showing 6 changed files with 69 additions and 5 deletions.
7 changes: 6 additions & 1 deletion playwright/_impl/_artifact.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

import pathlib
from pathlib import Path
from typing import Dict, Optional, Union, cast
from typing import AsyncIterator, Dict, Optional, Union, cast

from playwright._impl._connection import ChannelOwner, from_channel
from playwright._impl._helper import Error, make_dirs_for_file, patch_error_message
Expand All @@ -41,6 +41,11 @@ async def save_as(self, path: Union[str, Path]) -> None:
make_dirs_for_file(path)
await stream.save_as(path)

async def read_stream(self) -> AsyncIterator[bytes]:
stream = cast(Stream, from_channel(await self._channel.send("stream")))
async for chunk in stream.read_stream():
yield chunk

async def failure(self) -> Optional[str]:
return patch_error_message(await self._channel.send("failure"))

Expand Down
6 changes: 5 additions & 1 deletion playwright/_impl/_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

import pathlib
from pathlib import Path
from typing import TYPE_CHECKING, Optional, Union
from typing import TYPE_CHECKING, AsyncIterator, Optional, Union

from playwright._impl._artifact import Artifact

Expand Down Expand Up @@ -60,5 +60,9 @@ async def path(self) -> pathlib.Path:
async def save_as(self, path: Union[str, Path]) -> None:
await self._artifact.save_as(path)

async def read_stream(self) -> AsyncIterator[bytes]:
async for chunk in self._artifact.read_stream():
yield chunk

async def cancel(self) -> None:
return await self._artifact.cancel()
9 changes: 8 additions & 1 deletion playwright/_impl/_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

import base64
from pathlib import Path
from typing import Dict, Union
from typing import AsyncIterator, Dict, Union

from playwright._impl._connection import ChannelOwner

Expand All @@ -36,6 +36,13 @@ async def save_as(self, path: Union[str, Path]) -> None:
)
await self._loop.run_in_executor(None, lambda: file.close())

async def read_stream(self) -> AsyncIterator[bytes]:
while True:
binary = await self._channel.send("read", {"size": 1024 * 1024})
if not binary:
break
yield base64.b64decode(binary)

async def read_all(self) -> bytes:
binary = b""
while True:
Expand Down
11 changes: 10 additions & 1 deletion playwright/async_api/_generated.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@

import pathlib
import typing
from typing import Literal
from typing import AsyncIterator, Literal

from playwright._impl._accessibility import Accessibility as AccessibilityImpl
from playwright._impl._api_structures import (
Expand Down Expand Up @@ -6852,6 +6852,15 @@ async def save_as(self, path: typing.Union[str, pathlib.Path]) -> None:

return mapping.from_maybe_impl(await self._impl_obj.save_as(path=path))

async def read_stream(self) -> AsyncIterator[bytes]:
"""Download.read_stream

Yields a readable stream chunks for a successful download, or throws for a failed/canceled download.
"""

async for chunk in mapping.from_maybe_impl(self._impl_obj.read_stream()):
yield chunk

async def cancel(self) -> None:
"""Download.cancel

Expand Down
11 changes: 10 additions & 1 deletion playwright/sync_api/_generated.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@

import pathlib
import typing
from typing import Literal
from typing import Iterable, Literal

from playwright._impl._accessibility import Accessibility as AccessibilityImpl
from playwright._impl._api_structures import (
Expand Down Expand Up @@ -6962,6 +6962,15 @@ def save_as(self, path: typing.Union[str, pathlib.Path]) -> None:

return mapping.from_maybe_impl(self._sync(self._impl_obj.save_as(path=path)))

def read_stream(self) -> Iterable[bytes]:
"""Download.read_stream

Yields a readable stream chunks for a successful download, or throws for a failed/canceled download.
"""

for chunk in mapping.from_maybe_impl(self._sync(self._impl_obj.read_stream())):
yield chunk

def cancel(self) -> None:
"""Download.cancel

Expand Down
30 changes: 30 additions & 0 deletions tests/async/test_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,16 @@ def handle_download_with_file_name(request: TestServerRequest) -> None:
request.write(b"Hello world")
request.finish()

def handle_download_big_file(request: TestServerRequest) -> None:
request.setHeader("Content-Type", "application/octet-stream")
request.setHeader("Content-Disposition", "attachment")
request.write(b"A" * 1024 * 1024)
request.write(b"B")
request.finish()

server.set_route("/download", handle_download)
server.set_route("/downloadWithFilename", handle_download_with_file_name)
server.set_route("/downloadBigFile", handle_download_big_file)
yield


Expand Down Expand Up @@ -381,3 +389,25 @@ def handle_download(request: TestServerRequest) -> None:
await download.cancel()
assert await download.failure() == "canceled"
await page.close()


async def test_stream_reading(browser: Browser, server: Server) -> None:
page = await browser.new_page(accept_downloads=True)
await page.set_content(f'<a href="{server.PREFIX}/download">download</a>')
async with page.expect_download() as download_info:
await page.click("a")
download = await download_info.value
data = b"".join([chunk async for chunk in download.read_stream()])
assert data == b"Hello world"
await page.close()


async def test_stream_reading_multiple_chunks(browser: Browser, server: Server) -> None:
page = await browser.new_page(accept_downloads=True)
await page.set_content(f'<a href="{server.PREFIX}/downloadBigFile">download</a>')
async with page.expect_download() as download_info:
await page.click("a")
download = await download_info.value
data = b"".join([chunk async for chunk in download.read_stream()])
assert data == b"A" * 1024 * 1024 + b"B"
await page.close()

0 comments on commit e54eb53

Please sign in to comment.