fix: screenshot type inferred from path file extension (microsoft#2951)

Author: xianghuzhao

playwright/_impl/_element_handle.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import base64
+import mimetypes
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -323,6 +324,8 @@ async def screenshot(
     ) -> bytes:
         params = locals_to_params(locals())
         if "path" in params:
+            if "type" not in params:
+                params["type"] = determine_screenshot_type(params["path"])
             del params["path"]
         if "mask" in params:
             params["mask"] = list(
@@ -450,3 +453,12 @@ def convert_select_option_values(
         elements = list(map(lambda e: e._channel, element))
 
     return dict(options=options, elements=elements)
+
+
+def determine_screenshot_type(path: Union[str, Path]) -> Literal["jpeg", "png"]:
+    mime_type, _ = mimetypes.guess_type(path)
+    if mime_type == "image/png":
+        return "png"
+    if mime_type == "image/jpeg":
+        return "jpeg"
+    raise Error(f'Unsupported screenshot mime type for path "{path}": {mime_type}')

playwright/_impl/_page.py

@@ -51,7 +51,7 @@
 )
 from playwright._impl._console_message import ConsoleMessage
 from playwright._impl._download import Download
-from playwright._impl._element_handle import ElementHandle
+from playwright._impl._element_handle import ElementHandle, determine_screenshot_type
 from playwright._impl._errors import Error, TargetClosedError, is_target_closed_error
 from playwright._impl._event_context_manager import EventContextManagerImpl
 from playwright._impl._file_chooser import FileChooser
@@ -800,6 +800,8 @@ async def screenshot(
     ) -> bytes:
         params = locals_to_params(locals())
         if "path" in params:
+            if "type" not in params:
+                params["type"] = determine_screenshot_type(params["path"])
             del params["path"]
         if "mask" in params:
             params["mask"] = list(

tests/async/test_screenshot.py

@@ -12,13 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from pathlib import Path
 from typing import Callable
 
+from PIL import Image
+
 from playwright.async_api import Page
 from tests.server import Server
 from tests.utils import must
 
 
+def assert_image_file_format(path: Path, image_format: str) -> None:
+    with Image.open(path) as img:
+        assert img.format == image_format
+
+
 async def test_should_screenshot_with_mask(
     page: Page, server: Server, assert_to_be_golden: Callable[[bytes, str], None]
 ) -> None:
@@ -43,3 +51,29 @@ async def test_should_screenshot_with_mask(
         ),
         "mask-should-work-with-element-handle.png",
     )
+
+
+async def test_should_infer_screenshot_type_from_path(
+    page: Page, tmp_path: Path
+) -> None:
+    output_png_file = tmp_path / "foo.png"
+    await page.screenshot(path=output_png_file)
+    assert_image_file_format(output_png_file, "PNG")
+
+    output_jpeg_file = tmp_path / "bar.jpeg"
+    await page.screenshot(path=output_jpeg_file)
+    assert_image_file_format(output_jpeg_file, "JPEG")
+
+    output_jpg_file = tmp_path / "bar.jpg"
+    await page.screenshot(path=output_jpg_file)
+    assert_image_file_format(output_jpg_file, "JPEG")
+
+
+async def test_should_screenshot_with_type_argument(page: Page, tmp_path: Path) -> None:
+    output_jpeg_with_png_extension = tmp_path / "foo_jpeg.png"
+    await page.screenshot(path=output_jpeg_with_png_extension, type="jpeg")
+    assert_image_file_format(output_jpeg_with_png_extension, "JPEG")
+
+    output_png_with_jpeg_extension = tmp_path / "bar_png.jpeg"
+    await page.screenshot(path=output_png_with_jpeg_extension, type="png")
+    assert_image_file_format(output_png_with_jpeg_extension, "PNG")