robotframework · chriscallan · Apr 10, 2024 · Apr 10, 2024 · Apr 10, 2024 · Apr 11, 2024
diff --git a/atest/resources/atest_resource.robot b/atest/resources/atest_resource.robot
@@ -33,6 +33,7 @@ ${RUNNER DEFAULTS}
 ...               --ConsoleMarkers OFF
 ...               --PYTHONPATH "${CURDIR}${/}..${/}testresources${/}testlibs"
 ...               --PYTHONPATH "${CURDIR}${/}..${/}testresources${/}listeners"
+${SCREENSHOT DIR}     %{TEMPDIR}${/}robot_atest_screenshots
 
 *** Keywords ***
 Run Tests

diff --git a/atest/robot/standard_libraries/screenshot/take_screenshot.robot b/atest/robot/standard_libraries/screenshot/take_screenshot.robot
@@ -15,6 +15,14 @@ Basename May Be Defined
     ${tc}=  Check Test Case  ${TESTNAME}
     Check Embedding In Log  ${tc.kws[0].kws[0].msgs[1]}  foo_1.jpg
 
+Basename May Be Defined With Screenshot Format Of PNG
+    ${tc} =    Check Test Case    ${TESTNAME}
+    Check Embedding In Log    ${tc.kws[0].kws[0].msgs[1]}    foo_1.png
+
+Basename May Be Defined With Screenshot Format Of TIFF
+    ${tc} =    Check Test Case    ${TESTNAME}
+    Check Embedding In Log    ${tc.kws[0].kws[0].msgs[1]}    foo_1.tiff
+
 Basename With Extension Turns Off Index Generation
     ${tc}=  Check Test Case  ${TESTNAME}
     Check Embedding In Log  ${tc.kws[0].kws[0].msgs[1]}  xxx.jpg
@@ -38,9 +46,17 @@ Without Embedding
 *** Keywords ***
 Check Embedding In Log
     [Arguments]  ${message}  ${path}  ${width}=800px
-    Check Log Message  ${message}  <a href="${path}"><img src="${path}" width="${width}"></a>  HTML
+    ${rel_dir} =    Get Screenshot Dir As Relative Path    ${path}
+    Check Log Message  ${message}  <a href="${rel_dir}"><img src="${rel_dir}" width="${width}"></a>  HTML
 
 Check Linking In Log
     [Arguments]  ${message}  ${file}
-    ${path} =  Normalize Path  ${OUTDIR}/${file}
-    Check Log Message  ${message}  Screenshot saved to '<a href="${file}">${path}</a>'.  HTML
+    ${path} =  Set Variable    ${SCREENSHOT DIR}/${file}
+    ${rel_dir} =    Get Screenshot Dir As Relative Path        ${file}
+    Check Log Message  ${message}  Screenshot saved to '<a href="${rel_dir}">${path}</a>'.  HTML
+
+Get Screenshot Dir As Relative Path
+    [Arguments]    ${path}
+    ${parent_dir} =    Evaluate    "${SCREENSHOT DIR}".split('/')[-1]
+    ${ret_val} =       Set Variable    ../${parent_dir}/${path}
+    RETURN    ${ret_val}
diff --git a/atest/testdata/standard_libraries/screenshot/screenshot_resource.robot b/atest/testdata/standard_libraries/screenshot/screenshot_resource.robot
@@ -17,6 +17,7 @@ Save Start Time
     Set Test Variable  \${START TIME}
 
 Screenshots Should Exist
-    [Arguments]  ${directory}  @{files}
-    @{actual files}=  List Directory  ${directory}  *.jp*g
+    [Arguments]  ${directory}  @{files}    ${format}=jpg
+    ${file_ext_re} =    Set Variable If    "${format.lower()}" == "jpg"    *.jp*g    *.${format}
+    @{actual files}=  List Directory  ${directory}  ${file_ext_re}
     Lists Should Be Equal  ${actual files}  ${files}
diff --git a/atest/testdata/standard_libraries/screenshot/take_screenshot.robot b/atest/testdata/standard_libraries/screenshot/take_screenshot.robot
@@ -1,49 +1,98 @@
 *** Settings ***
-Suite Setup     Remove Files  ${OUTPUTDIR}/*.jp*g
-Test Setup      Save Start Time
-Test Teardown   Remove Files  ${OUTPUTDIR}/*.jp*g
+Library        OperatingSystem
+Suite Setup     Remove Files  ${OUTPUTDIR}/*.jp*g    ${OUTPUTDIR}/*.png    ${OUTPUTDIR}/*.tiff    ${OUTPUTDIR}/*.bmp
+Test Setup      Take Screenshot Test Setup
+Test Teardown   Take Screenshot Test Teardown
 Resource        screenshot_resource.robot
 
 *** Variables ***
 ${FIRST_SCREENSHOT}     screenshot_1.jpg
 ${SECOND_SCREENSHOT}    screenshot_2.jpg
+${SCREENSHOT DIR} =     %{TEMPDIR}${/}robot_atest_screenshots
+${BASENAME} =           ${SCREENSHOT DIR}${/}screenshot
+${FIRST_SCREENSHOT} =   ${BASENAME}_1.jpg
 
 *** Test Cases ***
 Screenshot Is Embedded in Log File
     ${path}=    Take Screenshot and Verify    ${FIRST_SCREENSHOT}
-    Should Be Equal    ${path}    ${OUTPUTDIR}${/}${FIRST_SCREENSHOT}
+    ${path_normalized} =    Normalize Path    ${SCREENSHOT DIR}
+    Should Be Equal    ${path}    ${SCREENSHOT DIR}${/}${FIRST_SCREENSHOT}
 
 Each Screenshot Gets Separate Index
     Take Screenshot and Verify    ${FIRST_SCREENSHOT}
     Take Screenshot and Verify    ${FIRST_SCREENSHOT}    ${SECOND_SCREENSHOT}
 
 Basename May Be Defined
     Repeat Keyword    2    Take Screenshot    foo
-    Screenshots Should Exist    ${OUTPUTDIR}    foo_1.jpg    foo_2.jpg
+    Screenshots Should Exist    ${SCREENSHOT DIR}    foo_1.jpg    foo_2.jpg
+
+Basename May Be Defined With Screenshot Format Of PNG
+    Repeat Keyword    2    Take Screenshot    foo    img_format=png
+    Screenshots Should Exist    ${SCREENSHOT DIR}    foo_1.png    foo_2.png    format=png
+
+Basename May Be Defined With Screenshot Format Of TIFF
+    Repeat Keyword    2    Take Screenshot    foo    img_format=tiff
+    Screenshots Should Exist    ${SCREENSHOT DIR}    foo_1.tiff    foo_2.tiff    format=tiff
 
 Basename With Extension Turns Off Index Generation
-    Repeat Keyword    3    Take Screenshot    xxx.jpg
-    Repeat Keyword    2    Take Screenshot    yyy.jpeg
-    Screenshots Should Exist  ${OUTPUTDIR}    xxx.jpg    yyy.jpeg
+    Repeat Keyword    3    Take Screenshot    xxx.jpg    img_format=jpg
+    Repeat Keyword    2    Take Screenshot    yyy.jpeg    img_format=jpg
+    Screenshots Should Exist  ${SCREENSHOT DIR}    xxx.jpg    yyy.jpeg    format=jpg
+    Repeat Keyword    3    Take Screenshot    xxx.png    img_format=png
+    Screenshots Should Exist  ${SCREENSHOT DIR}    xxx.png    format=png
+    Repeat Keyword    3    Take Screenshot    xxx.bmp    img_format=bmp
+    Screenshots Should Exist  ${SCREENSHOT DIR}    xxx.bmp    format=bmp
 
 Name as `pathlib.Path`
     Take Screenshot    ${{pathlib.Path('name.jpg')}}
-    Screenshots Should Exist    ${OUTPUTDIR}    name.jpg
+    Screenshots Should Exist    ${SCREENSHOT DIR}    name.jpg
+
+Name as `pathlib.Path` - Format PNG
+    Take Screenshot    ${{pathlib.Path('name.png')}}    img_format=png
+    Screenshots Should Exist    ${SCREENSHOT DIR}    name.png    format=png
+
+Name as `pathlib.Path` - Format TIFF
+    Take Screenshot    ${{pathlib.Path('name.tiff')}}    img_format=tiff
+    Screenshots Should Exist    ${SCREENSHOT DIR}    name.tiff    format=tiff
 
 Screenshot Width Can Be Given
     Take Screenshot    width=300px
-    Screenshots Should Exist    ${OUTPUTDIR}    ${FIRST_SCREENSHOT}
+    Screenshots Should Exist    ${SCREENSHOT DIR}    ${FIRST_SCREENSHOT}
+
+Screenshot Width Can Be Given For PNG
+    Take Screenshot    foo.png    width=300px    img_format=png
+    Screenshots Should Exist    ${SCREENSHOT DIR}    foo.png    format=png
+
+Screenshot Width Can Be Given For BMP
+    Take Screenshot    foo.bmp    width=300px    img_format=bmp
+    Screenshots Should Exist    ${SCREENSHOT DIR}    foo.bmp    format=bmp
 
 Basename With Non-existing Directories Fails
-    [Documentation]  FAIL Directory '${OUTPUTDIR}${/}non-existing' where to save the screenshot does not exist
-    Take Screenshot    ${OUTPUTDIR}${/}non-existing${/}foo
+    Run Keyword And Expect Error    Directory '/non-existing' where to save the screenshot does not exist
+    ...    Take Screenshot    ${/}non-existing${/}foo
 
 Without Embedding
     Take Screenshot Without Embedding    no_embed.jpeg
 
+Take Screenshot Without Embedding For PNG
+    Take Screenshot Without Embedding    no_embed.png    img_format=png
+
+Take Screenshot Without Embedding for TIFF
+    Take Screenshot Without Embedding    no_embed.tiff    img_format=tiff
+
 *** Keywords ***
+Take Screenshot Test Setup 
+    Save Start Time
+    Create Directory    ${SCREENSHOT DIR}
+    Set Screenshot Directory    ${SCREENSHOT DIR}
+
+Take Screenshot Test Teardown
+    Remove Directory    ${SCREENSHOT DIR}    recursive=${True}
+    Remove Files  ${SCREENSHOT DIR}/*.jp*g    ${SCREENSHOT DIR}/*.png
+    ...    ${SCREENSHOT DIR}/*.tiff    ${SCREENSHOT DIR}/*.bmp
+
 Take Screenshot And Verify
     [Arguments]    @{expected files}
     ${path}=    Take Screenshot
-    Screenshots Should Exist    ${OUTPUTDIR}    @{expected files}
+    Screenshots Should Exist    ${SCREENSHOT DIR}    @{expected files}
     RETURN    ${path}
diff --git a/src/robot/libraries/Screenshot.py b/src/robot/libraries/Screenshot.py
@@ -57,8 +57,10 @@ class Screenshot:
     one found will be used.
 
     - wxPython :: http://wxpython.org :: Generic Python GUI toolkit.
+        - supports: png, jpeg, gif, pcx, pnm, tiff, tga, ico, cur
     - PyGTK :: http://pygtk.org :: This module is available by default on most
       Linux distributions.
+        - supports:
     - Pillow :: http://python-pillow.github.io ::
       Only works on Windows. Also the original PIL package is supported.
     - Scrot :: http://en.wikipedia.org/wiki/Scrot :: Not used on Windows.
@@ -85,8 +87,26 @@ class Screenshot:
 
     ROBOT_LIBRARY_SCOPE = 'TEST SUITE'
     ROBOT_LIBRARY_VERSION = get_version()
-
-    def __init__(self, screenshot_directory=None, screenshot_module=None):
+    IMAGE_FORMATS = {
+        "jpg": {"extensions": ("jpg", "jpeg",),
+                "wx_bmp_type": "BITMAP_TYPE_JPEG",
+                "gtk_type": "jpeg",
+                "pil_type": "JPEG"},
+        "png": {"extensions": ("png",),
+                "wx_bmp_type": "BITMAP_TYPE_PNG",
+                "gtk_type": "png",
+                "pil_type": "PNG"},
+        "tiff": {"extensions": ("tiff", "tif",),
+                 "wx_bmp_type": "BITMAP_TYPE_TIFF",
+                 "gtk_type": None,
+                 "pil_type": "TIFF"},
+        "bmp": {"extensions": ("bmp",),
+                "wx_bmp_type": "BITMAP_TYPE_BMP",
+                "gtk_type": "bmp",
+                "pil_type": "BMP"}
+    }
+
+    def __init__(self, screenshot_directory=None, screenshot_module=None, image_format="jpg"):
         """Configure where screenshots are saved.
 
         If ``screenshot_directory`` is not given, screenshots are saved into
@@ -104,8 +124,9 @@ def __init__(self, screenshot_directory=None, screenshot_module=None):
         | Library   | Screenshot | ${TEMPDIR} |
         | Library   | Screenshot | screenshot_module=PyGTK |
         """
+        self._out_img_format = image_format
         self._given_screenshot_dir = self._norm_path(screenshot_directory)
-        self._screenshot_taker = ScreenshotTaker(screenshot_module)
+        self._screenshot_taker = ScreenshotTaker(screenshot_module, self._out_img_format)
 
     def _norm_path(self, path):
         if not path:
@@ -136,15 +157,17 @@ def set_screenshot_directory(self, path):
 
         The directory can also be set in `importing`.
         """
+        print(f"path passed in is: {path}")
         path = self._norm_path(path)
         if not os.path.isdir(path):
             raise RuntimeError("Directory '%s' does not exist." % path)
         old = self._screenshot_dir
         self._given_screenshot_dir = path
         return old
 
-    def take_screenshot(self, name="screenshot", width="800px"):
-        """Takes a screenshot in JPEG format and embeds it into the log file.
+    def take_screenshot(self, name="screenshot", width="800px", img_format="jpg"):
+        """Takes a screenshot in the specified format, jpg by default, and embeds
+        it into the log file
 
         Name of the file where the screenshot is stored is derived from the
         given ``name``. If the ``name`` ends with extension ``.jpg`` or
@@ -158,27 +181,33 @@ def take_screenshot(self, name="screenshot", width="800px"):
 
         ``width`` specifies the size of the screenshot in the log file.
 
+        ``img_format`` specifies the image format for the screenshot, default is "jpg"
+
         Examples: (LOGDIR is determined automatically by the library)
         | Take Screenshot |                  |     | # LOGDIR/screenshot_1.jpg (index automatically incremented) |
         | Take Screenshot | mypic            |     | # LOGDIR/mypic_1.jpg (index automatically incremented) |
         | Take Screenshot | ${TEMPDIR}/mypic |     | # /tmp/mypic_1.jpg (index automatically incremented) |
         | Take Screenshot | pic.jpg          |     | # LOGDIR/pic.jpg (always uses this file) |
         | Take Screenshot | images/login.jpg | 80% | # Specify both name and width. |
         | Take Screenshot | width=550px      |     | # Specify only width. |
+        | Take Screenshot | img_format=png   |     | # Specify only image format, png in this case |
+        | Take Screenshot | png_image        | img_format=tiff  | # Specify name and image format, "tiff" in this case |
 
         The path where the screenshot is saved is returned.
         """
+        self._out_img_format = img_format
         path = self._save_screenshot(name)
         self._embed_screenshot(path, width)
         return path
 
-    def take_screenshot_without_embedding(self, name="screenshot"):
+    def take_screenshot_without_embedding(self, name="screenshot", img_format="jpg"):
         """Takes a screenshot and links it from the log file.
 
         This keyword is otherwise identical to `Take Screenshot` but the saved
         screenshot is not embedded into the log file. The screenshot is linked
         so it is nevertheless easily available.
         """
+        self._out_img_format = img_format
         path = self._save_screenshot(name)
         self._link_screenshot(path)
         return path
@@ -208,12 +237,18 @@ def _validate_screenshot_path(self, path):
         return path
 
     def _get_screenshot_path(self, basename):
-        if basename.lower().endswith(('.jpg', '.jpeg')):
+        try:
+            # retrieve a tuple of the common file extension variations
+            file_extensions = Screenshot.IMAGE_FORMATS[self._out_img_format]["extensions"]
+        except KeyError as ke:
+            logger.error(f"KeyError in _get_screenshot_path(): {ke}")
+        if any([basename.lower().endswith(x) for x in file_extensions]):
             return os.path.join(self._screenshot_dir, basename)
         index = 0
         while True:
             index += 1
-            path = os.path.join(self._screenshot_dir, "%s_%d.jpg" % (basename, index))
+            # use the image format key for the file extension
+            path = os.path.join(self._screenshot_dir, "%s_%d.%s" % (basename, index, self._out_img_format))
             if not os.path.exists(path):
                 return path
 
@@ -230,7 +265,8 @@ def _link_screenshot(self, path):
 
 class ScreenshotTaker:
 
-    def __init__(self, module_name=None):
+    def __init__(self, module_name=None, image_format='jpg'):
+        self._image_format = image_format
         self._screenshot = self._get_screenshot_taker(module_name)
         self.module = self._screenshot.__name__.split('_')[1]
         self._wx_app_reference = None
@@ -289,7 +325,7 @@ def _get_default_screenshot_taker(self):
                 return screenshot_taker
 
     def _osx_screenshot(self, path):
-        if self._call('screencapture', '-t', 'jpg', path) != 0:
+        if self._call('screencapture', '-t', self._image_format, path) != 0:
             raise RuntimeError("Using 'screencapture' failed.")
 
     def _call(self, *command):
@@ -304,9 +340,13 @@ def _scrot(self):
         return os.sep == '/' and self._call('scrot', '--version') == 0
 
     def _scrot_screenshot(self, path):
-        if not path.endswith(('.jpg', '.jpeg')):
-            raise RuntimeError("Scrot requires extension to be '.jpg' or "
-                               "'.jpeg', got '%s'." % os.path.splitext(path)[1])
+        try:
+            file_extensions = Screenshot.IMAGE_FORMATS[self._image_format]["extensions"]
+        except KeyError as ke:
+            logger.error(f"KeyError in _scrot_screenshot(): {ke}")
+        if any([path.endswith(x) for x in file_extensions]):
+            raise RuntimeError("Scrot requires extension to be like %s, "
+                               "but got '%s'." % (str(file_extensions), os.path.splitext(path)[1]))
         if os.path.exists(path):
             os.remove(path)
         if self._call('scrot', '--silent', path) != 0:
@@ -325,7 +365,7 @@ def _wx_screenshot(self, path):
         memory.SelectObject(bitmap)
         memory.Blit(0, 0, width, height, context, -1, -1)
         memory.SelectObject(wx.NullBitmap)
-        bitmap.SaveFile(path, wx.BITMAP_TYPE_JPEG)
+        bitmap.SaveFile(path, Screenshot.IMAGE_FORMATS[self._image_format]["wx_bmp_type"])
 
     def _gtk_screenshot(self, path):
         window = gdk.get_default_root_window()
@@ -337,10 +377,10 @@ def _gtk_screenshot(self, path):
                                   0, 0, 0, 0, width, height)
         if not pb:
             raise RuntimeError('Taking screenshot failed.')
-        pb.save(path, 'jpeg')
+        pb.save(path, Screenshot.IMAGE_FORMATS[self._image_format]["gtk_type"])
 
     def _pil_screenshot(self, path):
-        ImageGrab.grab().save(path, 'JPEG')
+        ImageGrab.grab().save(path, Screenshot.IMAGE_FORMATS[self._image_format]["pil_type"])
 
     def _no_screenshot(self, path):
         raise RuntimeError('Taking screenshots is not supported on this platform '