In Gauge [python] custom_screenshot_writer not working when driver is none

[KrishnaAstra]

I use Gauge and have created the following Custom screenshot method to get full screen screenshots, It works fine.. However when the script fails before driver instantiation I had returned None but it fails...

Custom method:
@custom_screenshot_writer
def take_screenshot():
"""Custom Utility for Screenshot capture as per the called drivers

Returns:
    screenshot path
"""
screenshot_name = "screenshot-{0}.png".format(uuid1().int)

try:
    if (DriverMgr().driver or DriverMgr().driver != None:
        driver: webdriver= DriverMgr().driver
            
        screenshot_as_png = driver.get_screenshot_as_png()
        screenshot_file_name = os.path.join(os.getenv("gauge_screenshots_dir"), screenshot_name)        
        screenshot_file = open(screenshot_file_name, "wb")
        screenshot_file.write(screenshot_as_png)
        screenshot_file.close
        return os.path.basename(screenshot_file_name)
except Exception as exception:
        return None

Error Message: Exception calling application: expected str, bytes or os.PathLike object, not NoneType

I have some precondition and db verification steps which does not need driver to be present and where I dont have to have screesnhots as well. So what should I return so that Gauge does not throw me error

[KrishnaAstra]

Additional information
Python version: 3.11.4

Gauge version: 1.4.3
Commit Hash: f98dd40

Plugins

html-report (4.2.0)
python (0.3.17)
screenshot (0.1.0)
spectacle (0.1.4)

[zabil]

Have you tried modifying your code to return a blank screenshot when there's a exception?

from PIL import Image
import os
from uuid import uuid1
from selenium import webdriver

screenshot_name = "screenshot-{0}.png".format(uuid1().int)

try:
    if DriverMgr().driver or DriverMgr().driver != None:
        driver: webdriver = DriverMgr().driver
            
        screenshot_as_png = driver.get_screenshot_as_png()
        screenshot_file_name = os.path.join(os.getenv("gauge_screenshots_dir"), screenshot_name)        
        with open(screenshot_file_name, "wb") as screenshot_file:
            screenshot_file.write(screenshot_as_png)
        return os.path.basename(screenshot_file_name)
except Exception as exception:
    blank_screenshot = Image.new('RGB', (100, 100), color = (73, 109, 137))
    // Save blank screenshot or have a blank image presaved
    blank_screenshot_path = os.path.join(os.getenv("gauge_screenshots_dir"), "blank.png")
    blank_screenshot.save(blank_screenshot_path)
    return os.path.basename(blank_screenshot_path)

Please note this code was suggested by copilot so might want to test this.

[KrishnaAstra]

Hello Zabil,

Thank you for the suggestion..

This still creates a blank screenshot, even if i reduce the image size to 1,1 it still rearranges the HTML report like a placeholder for screenshot.[Refer below]

I was looking for a way to not add screenshot at all especially when it is not required or when script fails before Driver init happens.

Thank you.