Merge pull request #30 from Patolox/fix_file_extension

fix: Get file extension from file path and changes use of reserved word 'format' to 'fmt'

Author: thalissonvs

pydoll/browser/page.py

@@ -11,6 +11,7 @@
 from pydoll.commands.storage import StorageCommands
 from pydoll.connection.connection import ConnectionHandler
 from pydoll.element import WebElement
+from pydoll.exceptions import InvalidFileExtension
 from pydoll.mixins.find_elements import FindElementsMixin
 from pydoll.utils import decode_image_to_bytes
 
@@ -200,7 +201,13 @@ async def get_screenshot(self, path: str):
         Args:
             path (str): The file path to save the screenshot to.
         """
-        response = await self._execute_command(PageCommands.screenshot())
+        fmt = path.split('.')[-1]
+        if fmt not in {'jpeg', 'jpg', 'png'}:
+            raise InvalidFileExtension(f'{fmt} extension is not supported.')
+
+        response = await self._execute_command(
+            PageCommands.screenshot(fmt=fmt)
+        )
         screenshot_b64 = response['result']['data'].encode('utf-8')
         screenshot_bytes = decode_image_to_bytes(screenshot_b64)
         async with aiofiles.open(path, 'wb') as file:

pydoll/commands/page.py

@@ -71,13 +71,13 @@ def set_download_path(cls, path: str) -> dict:
 
     @classmethod
     def screenshot(
-        cls, format: str = 'jpeg', quality: int = 100, clip: dict = None
+        cls, fmt: str = 'jpeg', quality: int = 100, clip: dict = None
     ) -> dict:
         """
         Generates the command to capture a screenshot of the current page.
 
         Args:
-            format (str): The format of the image to be captured.
+            fmt (str): The format of the image to be captured.
                           Can be 'png' or 'jpeg'.
             quality (int): The quality of the image to be captured,
                            applicable only if the format is 'jpeg'.
@@ -89,7 +89,7 @@ def screenshot(
                   containing the method and parameters for the screenshot.
         """
         command = cls.SCREENSHOT_TEMPLATE.copy()
-        command['params']['format'] = format
+        command['params']['format'] = fmt
         command['params']['quality'] = quality
         if clip:
             command['params']['clip'] = clip

pydoll/element.py

@@ -189,7 +189,7 @@ async def get_screenshot(self, path: str):
             'scale': 1,
         }
         screenshot = await self._connection_handler.execute_command(
-            PageCommands.screenshot(format='jpeg', clip=clip)
+            PageCommands.screenshot(fmt='jpeg', clip=clip)
         )
         async with aiofiles.open(path, 'wb') as file:
             image_bytes = decode_image_to_bytes(screenshot['result']['data'])

pydoll/exceptions.py

@@ -80,3 +80,10 @@ class ElementNotInteractable(Exception):
 
     def __str__(self):
         return self.message
+
+
+class InvalidFileExtension(Exception):
+    message = 'The file extension provided is not supported'
+
+    def __str__(self):
+        return self.message

tests/test_page_commands.py

@@ -33,7 +33,7 @@ def test_screenshot_jpeg():
         },
     }
     assert (
-        PageCommands.screenshot(format='jpeg', quality=80) == expected_command
+        PageCommands.screenshot(fmt='jpeg', quality=80) == expected_command
     )
 
 
@@ -45,7 +45,7 @@ def test_screenshot_png():
             'quality': 100,
         },
     }
-    assert PageCommands.screenshot(format='png') == expected_command
+    assert PageCommands.screenshot(fmt='png') == expected_command
 
 
 def test_screenshot_with_clip():