refactor: keep take_screenshot consistent

Author: thalissonvs

pydoll/browser/tab.py

@@ -557,7 +557,7 @@ async def refresh(
 
     async def take_screenshot(
         self,
-        path: Optional[str] = None,
+        path: Optional[str | Path] = None,
         quality: int = 100,
         beyond_viewport: bool = False,
         as_base64: bool = False,
@@ -582,17 +582,28 @@ async def take_screenshot(
         if not path and not as_base64:
             raise MissingScreenshotPath()
 
-        output_extension = path.split('.')[-1] if path else ScreenshotFormat.PNG
+        if path and isinstance(path, str):
+            output_extension = path.split('.')[-1]
+        elif path and isinstance(path, Path):
+            output_extension = path.suffix.lstrip('.')
+        else:
+            output_extension = ScreenshotFormat.JPEG
+
+        # Normalize jpg to jpeg (CDP only accepts jpeg)
+        output_extension = output_extension.replace('jpg', 'jpeg') if output_extension == 'jpg' else output_extension
+
         if not ScreenshotFormat.has_value(output_extension):
             raise InvalidFileExtension(f'{output_extension} extension is not supported.')
 
+        output_format = ScreenshotFormat.get_value(output_extension)
+
         logger.info(
             f'Taking screenshot: path={path}, quality={quality}, '
             f'beyond_viewport={beyond_viewport}, as_base64={as_base64}'
         )
         response: CaptureScreenshotResponse = await self._execute_command(
             PageCommands.capture_screenshot(
-                format=ScreenshotFormat.get_value(output_extension),
+                format=output_format,
                 quality=quality,
                 capture_beyond_viewport=beyond_viewport,
             )
@@ -612,15 +623,15 @@ async def take_screenshot(
 
         if path:
             screenshot_bytes = decode_base64_to_bytes(screenshot_data)
-            async with aiofiles.open(path, 'wb') as file:
+            async with aiofiles.open(str(path), 'wb') as file:
                 await file.write(screenshot_bytes)
             logger.info(f'Screenshot saved to: {path}')
 
         return None
 
     async def print_to_pdf(
         self,
-        path: str,
+        path: Optional[str | Path] = None,
         landscape: bool = False,
         display_header_footer: bool = False,
         print_background: bool = True,
@@ -631,7 +642,7 @@ async def print_to_pdf(
         Generate PDF of current page.
 
         Args:
-            path: File path for PDF output.
+            path: File path for PDF output. Required if as_base64=False.
             landscape: Use landscape orientation.
             display_header_footer: Include header/footer.
             print_background: Include background graphics.
@@ -640,6 +651,9 @@ async def print_to_pdf(
 
         Returns:
             Base64 PDF data if as_base64=True, None otherwise.
+        
+        Raises:
+            ValueError: If path is not provided when as_base64=False.
         """
         logger.info(
             f'Generating PDF: path={path}, landscape={landscape}, '
@@ -659,6 +673,9 @@ async def print_to_pdf(
             logger.info('PDF generated and returned as base64')
             return pdf_data
 
+        if path is None:
+            raise ValueError("path is required when as_base64=False")
+
         pdf_bytes = decode_base64_to_bytes(pdf_data)
         async with aiofiles.open(path, 'wb') as file:
             await file.write(pdf_bytes)
@@ -826,7 +843,7 @@ async def continue_with_auth(
 
     @asynccontextmanager
     async def expect_file_chooser(
-        self, files: Union[str, Path, list[Union[str, Path]]]
+        self, files: str | Path | list[str | Path]
     ) -> AsyncGenerator[None, None]:
         """
         Context manager for automatic file upload handling.

pydoll/elements/web_element.py

@@ -3,6 +3,7 @@
 import asyncio
 import json
 import logging
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import aiofiles
@@ -24,6 +25,8 @@
     ElementNotFound,
     ElementNotInteractable,
     ElementNotVisible,
+    InvalidFileExtension,
+    MissingScreenshotPath,
     WaitElementTimeout,
 )
 from pydoll.protocol.input.types import (
@@ -226,12 +229,48 @@ async def get_siblings_elements(
         logger.debug(f'Siblings found: {len(siblings)}')
         return siblings
 
-    async def take_screenshot(self, path: str, quality: int = 100):
+    async def take_screenshot(
+        self,
+        path: Optional[str | Path] = None,
+        quality: int = 100,
+        as_base64: bool = False,
+    ) -> Optional[str]:
         """
         Capture screenshot of this element only.
 
         Automatically scrolls element into view before capturing.
+
+        Args:
+            path: File path for screenshot (extension determines format).
+            quality: Image quality 0-100 (default 100).
+            as_base64: Return as base64 string instead of saving file.
+
+        Returns:
+            Base64 screenshot data if as_base64=True, None otherwise.
+
+        Raises:
+            InvalidFileExtension: If file extension not supported.
+            MissingScreenshotPath: If path is None and as_base64 is False.
         """
+        if not path and not as_base64:
+            raise MissingScreenshotPath()
+
+        if path and isinstance(path, str):
+            output_extension = path.split('.')[-1]
+        elif path and isinstance(path, Path):
+            output_extension = path.suffix.lstrip('.')
+        else:
+            output_extension = ScreenshotFormat.JPEG
+
+        # Normalize jpg to jpeg (CDP only accepts jpeg)
+        if output_extension == 'jpg':
+            output_extension = 'jpeg'
+
+        if not ScreenshotFormat.has_value(output_extension):
+            raise InvalidFileExtension(f'{output_extension} extension is not supported.')
+
+        file_format = ScreenshotFormat.get_value(output_extension)
+
         bounds = await self.get_bounds_using_js()
         clip = Viewport(
             x=bounds['x'],
@@ -241,18 +280,29 @@ async def take_screenshot(self, path: str, quality: int = 100):
             scale=1,
         )
         logger.debug(
-            f'Taking element screenshot: path={path}, quality={quality}, '
+            f'Taking element screenshot: path={path}, quality={quality}, as_base64={as_base64}, '
             f'clip={{x: {clip["x"]}, y: {clip["y"]}, w: {clip["width"]}, h: {clip["height"]}}}'
         )
+
         screenshot: CaptureScreenshotResponse = await self._connection_handler.execute_command(
             PageCommands.capture_screenshot(
-                format=ScreenshotFormat.JPEG, clip=clip, quality=quality
+                format=file_format, clip=clip, quality=quality
             )
         )
-        async with aiofiles.open(path, 'wb') as file:
-            image_bytes = decode_base64_to_bytes(screenshot['result']['data'])
-            await file.write(image_bytes)
-        logger.info(f'Element screenshot saved: {path}')
+
+        screenshot_data = screenshot['result']['data']
+
+        if as_base64:
+            logger.info('Element screenshot captured and returned as base64')
+            return screenshot_data
+
+        if path:
+            image_bytes = decode_base64_to_bytes(screenshot_data)
+            async with aiofiles.open(str(path), 'wb') as file:
+                await file.write(image_bytes)
+            logger.info(f'Element screenshot saved: {path}')
+
+        return None
 
     def get_attribute(self, name: str) -> Optional[str]:
         """
@@ -417,7 +467,7 @@ async def insert_text(self, text: str):
         logger.info(f'Inserting text on element (length={len(text)})')
         await self._execute_command(InputCommands.insert_text(text))
 
-    async def set_input_files(self, files: list[str]):
+    async def set_input_files(self, files: str | Path | list[str | Path]):
         """
         Set file paths for file input element.
 
@@ -432,9 +482,10 @@ async def set_input_files(self, files: list[str]):
             or self._attributes.get('type', '').lower() != 'file'
         ):
             raise ElementNotAFileInput()
-        logger.info(f'Setting input files: count={len(files)}')
+        files_list = [str(file) for file in files] if isinstance(files, list) else [str(files)]
+        logger.info(f'Setting input files: count={len(files_list)}')
         await self._execute_command(
-            DomCommands.set_file_input_files(files=files, object_id=self._object_id)
+            DomCommands.set_file_input_files(files=files_list, object_id=self._object_id)
         )
 
     async def type_text(self, text: str, interval: float = 0.1):