Is it possible to take screenshots when the tests are over with TimeoutException?

Question

My code:

    @pytest.fixture()
    def setup(request):
        webdriver.FirefoxOptions().add_argument('--headless')
        driver = webdriver.Firefox(executable_path=GeckoDriverManager().install())
        driver.get(url)
        WebDriverWait(driver, 15).until(
            EC.visibility_of_element_located(
                (By.XPATH, "//*[contains(@class, 'e2e-login-form')]//input[@name='username']")
            )
        )
        driver.find_element(By.XPATH, "//*[contains(@class, 'e2e-login-form')]//input[@name='username']").send_keys(admin_login)
        driver.find_element(By.XPATH, "//*[contains(@class, 'e2e-login-form')]//input[@name='password']").send_keys(admin_password)
        driver.find_element(By.XPATH, "//*[contains(@class, 'e2e-login-form')]//button[@type='submit']").click()
        driver.implicitly_wait(20)
        driver.maximize_window()
        request.cls.driver = driver
        before_fail = request.session.testsfailed
        yield
        if request.session.testsfailed != before_fail:
            allure.attach(driver.get_screenshot_as_png(), name="Test failed", attachment_type=AttachmentType.PNG)
        driver.quit()

Thank you very much for any help