add fallback of playwright locator screenshot

dexplo · Jan 30, 2025 · 70b973d · 70b973d
1 parent e9109b9
commit 70b973d
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/dataframe_image/converter/browser/playwright_converter.py b/dataframe_image/converter/browser/playwright_converter.py
@@ -82,7 +82,11 @@ def screenshot(self, html):
                     )
                     pass
                 page.wait_for_timeout(200)
-            screenshot_bytes = locator.screenshot(timeout=1000)
+            try:
+                screenshot_bytes = locator.screenshot(timeout=1000)
+            except Error:
+                logger.warning("Locator screenshot timed out. Taking full page screenshot instead.")
+                screenshot_bytes = page.screenshot(timeout=1000)
         im = Image.open(BytesIO(screenshot_bytes))
         return im
 
@@ -168,6 +172,10 @@ async def screenshot(self, html):
                     )
                     pass
                 page.wait_for_timeout(200)
-            screenshot_bytes = await locator.screenshot(timeout=1000)
+            try:
+                screenshot_bytes = await locator.screenshot(timeout=1000)
+            except Error:
+                logger.warning("Locator screenshot timed out. Taking full page screenshot instead.")
+                screenshot_bytes = await page.screenshot(timeout=1000)
         im = Image.open(BytesIO(screenshot_bytes))
         return im