Add save_screenshot()

Author: mdmintz

seleniumbase/fixtures/base_case.py

@@ -214,6 +214,9 @@ def wait_for_and_dismiss_alert(self, timeout=settings.LARGE_TIMEOUT):
     def wait_for_and_switch_to_alert(self, timeout=settings.LARGE_TIMEOUT):
         return page_actions.wait_for_and_switch_to_alert(self.driver, timeout)
 
+    def save_screenshot(self, name, folder=None):
+        return page_actions.save_screenshot(self.driver, name, folder)
+
 # PyTest-Specific Code #
 
     def setUp(self):

seleniumbase/fixtures/page_actions.py

@@ -19,6 +19,7 @@
 By.PARTIAL_LINK_TEXT
 """
 
+import os
 import time
 from seleniumbase.config import settings
 from selenium.webdriver.common.by import By
@@ -276,6 +277,22 @@ def find_visible_elements(driver, selector, by=By.CSS_SELECTOR):
     return [element for element in elements if element.is_displayed()]
 
 
+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.
+    """
+    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)
+    else:
+        screenshot_file = name
+    driver.get_screenshot_as_file(screenshot_file)
+
+
 def wait_for_ready_state_complete(driver, timeout=settings.EXTREME_TIMEOUT):
     """
     The DOM (Document Object Model) has a property called "readyState".