Change the algorithm for taking screenshots

Author: mdmintz

seleniumbase/core/log_helper.py

@@ -11,9 +11,15 @@ def log_screenshot(test_logpath, driver):
     screenshot_name = settings.SCREENSHOT_NAME
     screenshot_path = "%s/%s" % (test_logpath, screenshot_name)
     try:
-        driver.get_screenshot_as_file(screenshot_path)
+        element = driver.find_element_by_tag_name('body')
+        element_png = element.screenshot_as_png
+        with open(screenshot_path, "wb") as file:
+            file.write(element_png)
     except Exception:
-        print("WARNING: Unable to get screenshot for failure logs!")
+        try:
+            driver.get_screenshot_as_file(screenshot_path)
+        except Exception:
+            print("WARNING: Unable to get screenshot for failure logs!")
 
 
 def log_test_failure_data(test, test_logpath, driver, browser):

seleniumbase/fixtures/base_case.py

@@ -2449,6 +2449,7 @@ def switch_to_default_window(self):
         self.switch_to_window(0)
 
     def save_screenshot(self, name, folder=None):
+        """ The screenshot will be in PNG format. """
         return page_actions.save_screenshot(self.driver, name, folder)
 
     def get_new_driver(self, browser=None, headless=None,

seleniumbase/fixtures/page_actions.py

@@ -369,16 +369,25 @@ def save_screenshot(driver, name, folder=None):
     """
     Saves a screenshot to the current directory (or to a subfolder if provided)
     If the folder provided doesn't exist, it will get created.
+    The screenshot will be in PNG format.
     """
+    if "." not in name:
+        name = name + ".png"
     if folder:
         abs_path = os.path.abspath('.')
         file_path = abs_path + "/%s" % folder
         if not os.path.exists(file_path):
             os.makedirs(file_path)
-        screenshot_file = "%s/%s" % (file_path, name)
+        screenshot_path = "%s/%s" % (file_path, name)
     else:
-        screenshot_file = name
-    driver.get_screenshot_as_file(screenshot_file)
+        screenshot_path = name
+    try:
+        element = driver.find_element_by_tag_name('body')
+        element_png = element.screenshot_as_png
+        with open(screenshot_path, "wb") as file:
+            file.write(element_png)
+    except Exception:
+        driver.get_screenshot_as_file(screenshot_path)
 
 
 def _get_last_page(driver):