Merge pull request #741 from seleniumbase/text-overlay-on-images

Add the option to include overlay text on saved images

Author: mdmintz

docs/requirements.txt

@@ -1,5 +1,5 @@
 regex>=2020.11.13
-tqdm>=4.53.0
+tqdm>=4.54.0
 livereload==2.6.3;python_version>="3.6"
 Markdown==3.3.3
 readme-renderer==28.0

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)
 

requirements.txt

@@ -16,7 +16,7 @@ chardet==3.0.4
 urllib3==1.26.2
 requests==2.25.0
 selenium==3.141.0
-msedge-selenium-tools==3.141.2
+msedge-selenium-tools==3.141.3
 more-itertools==5.0.0;python_version<"3.5"
 more-itertools==8.6.0;python_version>="3.5"
 cssselect==1.1.0
@@ -64,6 +64,9 @@ pyotp==2.4.1
 boto==2.49.0
 cffi==1.14.4
 toml==0.10.2
+Pillow==6.2.2;python_version<"3.5"
+Pillow==7.2.0;python_version>="3.5" and python_version<"3.6"
+Pillow==8.0.1;python_version>="3.6"
 rich==9.2.0;python_version>="3.6" and python_version<"4.0"
 zipp==1.2.0;python_version<"3.6"
 zipp==3.4.0;python_version>="3.6"
@@ -73,6 +76,6 @@ pyflakes==2.1.1;python_version<"3.5"
 pyflakes==2.2.0;python_version>="3.5"
 tornado==5.1.1;python_version<"3.5"
 tornado==6.1;python_version>="3.5"
-allure-pytest==2.8.21
+allure-pytest==2.8.22
 pdfminer.six==20191110;python_version<"3.5"
 pdfminer.six==20201018;python_version>="3.5"

seleniumbase/__version__.py

@@ -1,2 +1,2 @@
 # seleniumbase package
-__version__ = "1.50.13"
+__version__ = "1.50.14"

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.

setup.py

@@ -49,7 +49,7 @@
         print("\n*** Installing twine: *** (Required for PyPI uploads)\n")
         os.system("python -m pip install --upgrade 'twine>=1.15.0'")
         print("\n*** Installing tqdm: *** (Required for PyPI uploads)\n")
-        os.system("python -m pip install --upgrade 'tqdm>=4.53.0'")
+        os.system("python -m pip install --upgrade 'tqdm>=4.54.0'")
         print("\n*** Publishing The Release to PyPI: ***\n")
         os.system('python -m twine upload dist/*')  # Requires ~/.pypirc Keys
         print("\n*** The Release was PUBLISHED SUCCESSFULLY to PyPI! :) ***\n")
@@ -120,7 +120,7 @@
         'urllib3==1.26.2',  # Must stay in sync with "requests"
         'requests==2.25.0',
         'selenium==3.141.0',
-        'msedge-selenium-tools==3.141.2',
+        'msedge-selenium-tools==3.141.3',
         'more-itertools==5.0.0;python_version<"3.5"',
         'more-itertools==8.6.0;python_version>="3.5"',
         'cssselect==1.1.0',
@@ -168,6 +168,9 @@
         'boto==2.49.0',
         'cffi==1.14.4',
         'toml==0.10.2',
+        'Pillow==6.2.2;python_version<"3.5"',
+        'Pillow==7.2.0;python_version>="3.5" and python_version<"3.6"',
+        'Pillow==8.0.1;python_version>="3.6"',
         'rich==9.2.0;python_version>="3.6" and python_version<"4.0"',
         'zipp==1.2.0;python_version<"3.6"',
         'zipp==3.4.0;python_version>="3.6"',
@@ -177,7 +180,7 @@
         'pyflakes==2.2.0;python_version>="3.5"',
         'tornado==5.1.1;python_version<"3.5"',
         'tornado==6.1;python_version>="3.5"',
-        'allure-pytest==2.8.21',
+        'allure-pytest==2.8.22',
         'pdfminer.six==20191110;python_version<"3.5"',
         'pdfminer.six==20201018;python_version>="3.5"',
     ],