code: def save_screenshot(driver, name):
      
      if hasattr(driver, 'save_screenshot'):
          screenshot_dir = os.environ.get('SCREENSHOT_DIR')
          if not screenshot_dir:
              LOGGER.warning('The SCREENSHOT_DIR environment variable was not set; not saving a screenshot')
              return
          elif not os.path.exists(screenshot_dir):
              os.makedirs(screenshot_dir)
          image_name = os.path.join(screenshot_dir, name + '.png')
          driver.save_screenshot(image_name)
  
      else:
          msg = (
              u"Browser does not support screenshots. "
              u"Could not save screenshot '{name}'"
          ).format(name=name)
  
          LOGGER.warning(msg)
docs: Save a screenshot of the browser.
  
      The location of the screenshot can be configured
      by the environment variable `SCREENSHOT_DIR`.  If not set,
      this defaults to the current working directory.
  
      Args:
          driver (selenium.webdriver): The Selenium-controlled browser.
          name (str): A name for the screenshot, which will be used in the output file name.
  
      Returns:
          None