Skip to content

Commit 254aabd

Browse files
authored
fix: screenshot type inferred from path file extension (#2951)
1 parent 1be34f2 commit 254aabd

File tree

3 files changed

+49
-1
lines changed

3 files changed

+49
-1
lines changed

playwright/_impl/_element_handle.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
import base64
16+
import mimetypes
1617
from pathlib import Path
1718
from typing import (
1819
TYPE_CHECKING,
@@ -323,6 +324,8 @@ async def screenshot(
323324
) -> bytes:
324325
params = locals_to_params(locals())
325326
if "path" in params:
327+
if "type" not in params:
328+
params["type"] = determine_screenshot_type(params["path"])
326329
del params["path"]
327330
if "mask" in params:
328331
params["mask"] = list(
@@ -450,3 +453,12 @@ def convert_select_option_values(
450453
elements = list(map(lambda e: e._channel, element))
451454

452455
return dict(options=options, elements=elements)
456+
457+
458+
def determine_screenshot_type(path: Union[str, Path]) -> Literal["jpeg", "png"]:
459+
mime_type, _ = mimetypes.guess_type(path)
460+
if mime_type == "image/png":
461+
return "png"
462+
if mime_type == "image/jpeg":
463+
return "jpeg"
464+
raise Error(f'Unsupported screenshot mime type for path "{path}": {mime_type}')

playwright/_impl/_page.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151
)
5252
from playwright._impl._console_message import ConsoleMessage
5353
from playwright._impl._download import Download
54-
from playwright._impl._element_handle import ElementHandle
54+
from playwright._impl._element_handle import ElementHandle, determine_screenshot_type
5555
from playwright._impl._errors import Error, TargetClosedError, is_target_closed_error
5656
from playwright._impl._event_context_manager import EventContextManagerImpl
5757
from playwright._impl._file_chooser import FileChooser
@@ -800,6 +800,8 @@ async def screenshot(
800800
) -> bytes:
801801
params = locals_to_params(locals())
802802
if "path" in params:
803+
if "type" not in params:
804+
params["type"] = determine_screenshot_type(params["path"])
803805
del params["path"]
804806
if "mask" in params:
805807
params["mask"] = list(

tests/async/test_screenshot.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,21 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from pathlib import Path
1516
from typing import Callable
1617

18+
from PIL import Image
19+
1720
from playwright.async_api import Page
1821
from tests.server import Server
1922
from tests.utils import must
2023

2124

25+
def assert_image_file_format(path: Path, image_format: str) -> None:
26+
with Image.open(path) as img:
27+
assert img.format == image_format
28+
29+
2230
async def test_should_screenshot_with_mask(
2331
page: Page, server: Server, assert_to_be_golden: Callable[[bytes, str], None]
2432
) -> None:
@@ -43,3 +51,29 @@ async def test_should_screenshot_with_mask(
4351
),
4452
"mask-should-work-with-element-handle.png",
4553
)
54+
55+
56+
async def test_should_infer_screenshot_type_from_path(
57+
page: Page, tmp_path: Path
58+
) -> None:
59+
output_png_file = tmp_path / "foo.png"
60+
await page.screenshot(path=output_png_file)
61+
assert_image_file_format(output_png_file, "PNG")
62+
63+
output_jpeg_file = tmp_path / "bar.jpeg"
64+
await page.screenshot(path=output_jpeg_file)
65+
assert_image_file_format(output_jpeg_file, "JPEG")
66+
67+
output_jpg_file = tmp_path / "bar.jpg"
68+
await page.screenshot(path=output_jpg_file)
69+
assert_image_file_format(output_jpg_file, "JPEG")
70+
71+
72+
async def test_should_screenshot_with_type_argument(page: Page, tmp_path: Path) -> None:
73+
output_jpeg_with_png_extension = tmp_path / "foo_jpeg.png"
74+
await page.screenshot(path=output_jpeg_with_png_extension, type="jpeg")
75+
assert_image_file_format(output_jpeg_with_png_extension, "JPEG")
76+
77+
output_png_with_jpeg_extension = tmp_path / "bar_png.jpeg"
78+
await page.screenshot(path=output_png_with_jpeg_extension, type="png")
79+
assert_image_file_format(output_png_with_jpeg_extension, "PNG")

0 commit comments

Comments
 (0)