Merge pull request #312 from seleniumbase/localized-screenshots

Add ability to save page elements as screenshot files

Author: mdmintz

.gitignore

@@ -70,6 +70,9 @@ report.xml
 # Tours
 tours_exported
 
+# Images
+images_exported
+
 # Automated Visual Testing
 visual_baseline
 

examples/image_test.py

@@ -0,0 +1,16 @@
+"""
+Pull an image from a website and save it as a PNG file.
+"""
+
+from seleniumbase import BaseCase
+
+
+class ImageTest(BaseCase):
+
+    def test_pull_image_from_website(self):
+        self.open("https://xkcd.com/1117/")
+        selector = "#comic"
+        file_name = "comic.png"
+        folder = "images_exported"
+        self.save_element_as_image_file(selector, file_name, folder)
+        print('"%s/%s" has been saved!' % (folder, file_name))

help_docs/method_summary.md

@@ -182,6 +182,10 @@ self.print_unique_links_with_status_codes()
 
 self.safe_execute_script(script)
 
+self.create_folder(folder)
+
+self.save_element_as_image_file(selector, file_name, folder=None)
+
 self.download_file(file_url, destination_folder=None)
 
 self.save_file_as(file_url, new_file_name, destination_folder=None)

seleniumbase/fixtures/base_case.py

@@ -1673,6 +1673,39 @@ def safe_execute_script(self, script):
             self.activate_jquery()  # It's a good thing we can define it here
             self.execute_script(script)
 
+    def create_folder(self, folder):
+        """ Creates a folder of the given name if it doesn't already exist. """
+        if folder.endswith("/"):
+            folder = folder[:-1]
+        if len(folder) < 1:
+            raise Exception("Minimum folder name length = 1.")
+        if not os.path.exists(folder):
+            try:
+                os.makedirs(folder)
+            except Exception:
+                pass
+
+    def save_element_as_image_file(self, selector, file_name, folder=None):
+        """ Take a screenshot of an element and save it as an image file.
+            If no folder is specified, will save it to the current folder. """
+        element = self.find_element(selector)
+        element_png = element.screenshot_as_png
+        if len(file_name.split('.')[0]) < 1:
+            raise Exception("Error: file_name length must be > 0.")
+        if not file_name.endswith(".png"):
+            file_name = file_name + ".png"
+        image_file_path = None
+        if folder:
+            if folder.endswith("/"):
+                folder = folder[:-1]
+            if len(folder) > 0:
+                self.create_folder(folder)
+                image_file_path = "%s/%s" % (folder, file_name)
+        if not image_file_path:
+            image_file_path = file_name
+        with open(image_file_path, "wb") as file:
+            file.write(element_png)
+
     def download_file(self, file_url, destination_folder=None):
         """ Downloads the file from the url to the destination folder.
             If no destination folder is specified, the default one is used.

setup.py

@@ -17,7 +17,7 @@
 
 setup(
     name='seleniumbase',
-    version='1.22.8',
+    version='1.22.9',
     description='Reliable Browser Automation & Testing Framework',
     long_description=long_description,
     long_description_content_type='text/markdown',