Add "take_screenshot" method for page and element (#102)

* Add "take_screenshot" for page and element

it returns the screenshot as base64 string instead of saving to local file

* Update CHANGELOG.md

---------

Co-authored-by: Stephan Lensky <8302875+stephanlensky@users.noreply.github.com>

Author: falmar

CHANGELOG.md

@@ -11,8 +11,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added `tab.take_screenshot` and `element.take_screenshot` methods to return screenshot as base64 string @falmar
+
 ### Changed
 
+- Changed `take.save_screenshot` and `element.save_screenshot` to use `self.take_screenshot` to fetch screenshot data and only perform the file saving @falmar
+
 ### Removed
 
 ## [0.6.1] - 2025-04-25

zendriver/core/element.py

@@ -810,6 +810,46 @@ async def query_selector(self, selector):
         await self.update()
         return await self.tab.query_selector(selector, self)
 
+    async def take_screenshot(
+        self,
+        format: str = "jpeg",
+        scale: typing.Optional[typing.Union[int, float]] = 1,
+    ):
+        """
+        Takes a screenshot of this element (only)
+        This is not the same as :py:obj:`Tab.take_screenshot`, which takes a "regular" screenshot
+
+        When the element is hidden, or has no size, or is otherwise not capturable, a RuntimeError is raised
+
+        :param format: jpeg or png (defaults to jpeg)
+        :type format: str
+        :param scale: the scale of the screenshot, eg: 1 = size as is, 2 = double, 0.5 is half
+        :return: screenshot data as base64 encoded
+        :rtype: str
+        """
+        pos = await self.get_position()
+        if not pos:
+            raise RuntimeError(
+                "could not determine position of element. probably because it's not in view, or hidden"
+            )
+        viewport = pos.to_viewport(scale)
+        await self.tab.sleep()
+
+        data = await self._tab.send(
+            cdp.page.capture_screenshot(
+                format, clip=viewport, capture_beyond_viewport=True
+            )
+        )
+
+        if not data:
+            from .connection import ProtocolException
+
+            raise ProtocolException(
+                "could not take screenshot. most possible cause is the page has not finished loading yet."
+            )
+
+        return str(data)
+
     async def save_screenshot(
         self,
         filename: typing.Optional[PathLike] = "auto",
@@ -830,14 +870,8 @@ async def save_screenshot(
         :return: the path/filename of saved screenshot
         :rtype: str
         """
-        pos = await self.get_position()
-        if not pos:
-            raise RuntimeError(
-                "could not determine position of element. probably because it's not in view, or hidden"
-            )
-        viewport = pos.to_viewport(scale)
-        path = None
         await self.tab.sleep()
+
         if not filename or filename == "auto":
             parsed = urllib.parse.urlparse(self.tab.target.url)
             parts = parsed.path.split("/")
@@ -848,26 +882,15 @@ async def save_screenshot(
             ext = ""
             if format.lower() in ["jpg", "jpeg"]:
                 ext = ".jpg"
-                format = "jpeg"
             elif format.lower() in ["png"]:
                 ext = ".png"
-                format = "png"
             path = pathlib.Path(candidate + ext)
         else:
             path = pathlib.Path(filename)
 
         path.parent.mkdir(parents=True, exist_ok=True)
-        data = await self._tab.send(
-            cdp.page.capture_screenshot(
-                format, clip=viewport, capture_beyond_viewport=True
-            )
-        )
-        if not data:
-            from .connection import ProtocolException
 
-            raise ProtocolException(
-                "could not take screenshot. most possible cause is the page has not finished loading yet."
-            )
+        data = await self.take_screenshot(format, scale)
 
         data_bytes = base64.b64decode(data)
         if not path:

zendriver/core/tab.py

@@ -1285,7 +1285,7 @@ async def download_file(self, url: str, filename: Optional[PathLike] = None):
                 arguments=[cdp.runtime.CallArgument(object_id=body.object_id)],
             )
         )
-
+      
     async def save_snapshot(self, filename: str = "snapshot.mhtml") -> None:
         """
         Saves a snapshot of the page.
@@ -1302,6 +1302,44 @@ async def save_snapshot(self, filename: str = "snapshot.mhtml") -> None:
 
         with open(filename, "w") as file:
             file.write(data)
+            
+    async def take_screenshot(
+        self,
+        format: str = "jpeg",
+        full_page: bool = False,
+    ) -> str:
+        """
+        Takes a screenshot of the page.
+        This is not the same as :py:obj:`Element.take_screenshot`, which takes a screenshot of a single element only
+
+        :param format: jpeg or png (defaults to jpeg)
+        :type format: str
+        :param full_page: when False (default) it captures the current viewport. when True, it captures the entire page
+        :type full_page: bool
+        :return: screenshot data as base64 encoded
+        :rtype: str
+        """
+        if self.target is None:
+            raise ValueError("target is none")
+
+        await self.sleep()  # update the target's url
+
+        if format.lower() in ["jpg", "jpeg"]:
+            format = "jpeg"
+        elif format.lower() in ["png"]:
+            format = "png"
+
+        data = await self.send(
+            cdp.page.capture_screenshot(
+                format_=format, capture_beyond_viewport=full_page
+            )
+        )
+        if not data:
+            raise ProtocolException(
+                "could not take screenshot. most possible cause is the page has not finished loading yet."
+            )
+
+        return str(data)
 
     async def save_screenshot(
         self,
@@ -1322,19 +1360,11 @@ async def save_screenshot(
         :return: the path/filename of saved screenshot
         :rtype: str
         """
-        if self.target is None:
-            raise ValueError("target is none")
-
-        await self.sleep()  # update the target's url
-        path = None
-
         if format.lower() in ["jpg", "jpeg"]:
             ext = ".jpg"
-            format = "jpeg"
 
         elif format.lower() in ["png"]:
             ext = ".png"
-            format = "png"
 
         if not filename or filename == "auto":
             parsed = urllib.parse.urlparse(self.target.url)
@@ -1347,15 +1377,9 @@ async def save_screenshot(
         else:
             path = pathlib.Path(filename)
         path.parent.mkdir(parents=True, exist_ok=True)
-        data = await self.send(
-            cdp.page.capture_screenshot(
-                format_=format, capture_beyond_viewport=full_page
-            )
-        )
-        if not data:
-            raise ProtocolException(
-                "could not take screenshot. most possible cause is the page has not finished loading yet."
-            )
+
+        data = await self.take_screenshot(format=format, full_page=full_page)
+
         import base64
 
         data_bytes = base64.b64decode(data)