selenium 错误时，截图保存

from selenium import webdriver
import time
driver = webdriver.Firefox()
driver.get("https://www.baidu.com")

try:
driver.find_element_by_id("kw_error").send_keys("selenium")
driver.find_element_by_id("su").click()
except:
driver.get_screenshot_as_file("D:\error01.png")
driver.quit()