Add the option to include overlay text on saved images

Author: mdmintz

examples/image_test.py

@@ -1,16 +1,51 @@
-"""
-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):
+        """ Pull an image from a website and save it as a PNG file. """
         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))
+
+    def test_add_text_overlay_to_image(self):
+        """ Add a text overlay to an image. """
+        self.open("https://xkcd.com/1117/")
+        selector = "#comic"
+        file_name = "image_overlay.png"
+        folder = "images_exported"
+        overlay_text = 'This is an XKCD comic!\nTitle: "My Sky"'
+        self.save_element_as_image_file(
+            selector, file_name, folder, overlay_text)
+        print('"%s/%s" has been saved!' % (folder, file_name))
+
+    def test_add_text_overlay_to_page_section(self):
+        """ Add a text overlay to a section of a page. """
+        self.open("https://xkcd.com/2200/")
+        selector = "#middleContainer"
+        file_name = "section_overlay.png"
+        folder = "images_exported"
+        overlay_text = (
+            'Welcome to %s\n'
+            'This is a comment added to the image.\n'
+            'Unreachable states come from logic errors.'
+            % self.get_current_url())
+        self.save_element_as_image_file(
+            selector, file_name, folder, overlay_text)
+        print('"%s/%s" has been saved!' % (folder, file_name))
+
+    def test_add_text_overlay_to_full_page(self):
+        """ Add a text overlay to a full page. """
+        self.open("https://xkcd.com/1922/")
+        self.remove_element("#bottom")
+        selector = "body"
+        file_name = "page_overlay.png"
+        folder = "images_exported"
+        overlay_text = ("A text overlay on %s" % self.get_current_url())
+        self.save_element_as_image_file(
+            selector, file_name, folder, overlay_text)
+        print('"%s/%s" has been saved!' % (folder, file_name))

help_docs/method_summary.md

@@ -277,7 +277,7 @@ self.create_folder(folder)
 
 self.choose_file(selector, file_path, by=By.CSS_SELECTOR, timeout=None)
 
-self.save_element_as_image_file(selector, file_name, folder=None)
+self.save_element_as_image_file(selector, file_name, folder=None, overlay_text="")
 
 self.download_file(file_url, destination_folder=None)
 

seleniumbase/fixtures/base_case.py

@@ -2784,9 +2784,11 @@ def choose_file(self, selector, file_path, by=By.CSS_SELECTOR,
         abs_path = os.path.abspath(file_path)
         self.add_text(selector, abs_path, by=by, timeout=timeout)
 
-    def save_element_as_image_file(self, selector, file_name, folder=None):
+    def save_element_as_image_file(
+            self, selector, file_name, folder=None, overlay_text=""):
         """ 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. """
+            If no folder is specified, will save it to the current folder.
+            If overlay_text is provided, will add that to the saved image. """
         element = self.wait_for_element_visible(selector)
         element_png = element.screenshot_as_png
         if len(file_name.split('.')[0]) < 1:
@@ -2804,6 +2806,26 @@ def save_element_as_image_file(self, selector, file_name, folder=None):
             image_file_path = file_name
         with open(image_file_path, "wb") as file:
             file.write(element_png)
+        # Add a text overlay if given
+        if type(overlay_text) is str and len(overlay_text) > 0:
+            from PIL import Image, ImageDraw
+            text_rows = overlay_text.split("\n")
+            len_text_rows = len(text_rows)
+            max_width = 0
+            for text_row in text_rows:
+                if len(text_row) > max_width:
+                    max_width = len(text_row)
+            image = Image.open(image_file_path)
+            draw = ImageDraw.Draw(image)
+            draw.rectangle(
+                (0, 0, (max_width * 6) + 6, 16 * len_text_rows),
+                fill=(236, 236, 28))
+            draw.text(
+                (4, 2),  # Coordinates
+                overlay_text,  # Text
+                (8, 38, 176)  # Color
+            )
+            image.save(image_file_path, 'PNG', quality=100, optimize=True)
 
     def download_file(self, file_url, destination_folder=None):
         """ Downloads the file from the url to the destination folder.