Streamlit as UI for Selenium scripts

app.py

import atexit
import streamlit as st
from selenium import webdriver
from selenium.webdriver.support.ui import WebDriverWait
from selenium.common.exceptions import WebDriverException


def init_chrome():
    chrome_options = webdriver.ChromeOptions()
    # chrome_options.add_argument("--headless=new")
    chrome_options.add_argument("--disable-extensions")
    return webdriver.Chrome(options=chrome_options)


# initialize selenium only once
if "driver" not in st.session_state:
    st.session_state.driver = init_chrome()

try:
    # attempt to access the first tab. will error out if the window or tab is closed
    first_tab = st.session_state.driver.window_handles[0]
except:
    # if the first tab is not accessible, restart chrome
    st.session_state.driver.quit()
    st.session_state.driver = init_chrome()

# get driver from session, init wait
driver = st.session_state.driver
wait = WebDriverWait(driver, 20)


# register an exit handler, quit driver when script stops
if "atexit_registered" not in st.session_state:

    def cleanup():
        print("cleanup")
        driver.quit()

    if not getattr(st.session_state, "atexit_registered", False):
        atexit.register(cleanup)
        st.session_state.atexit_registered = True


# always go first tab, prevents opening multiple windows
first_tab = driver.window_handles[0]
if driver.current_window_handle != first_tab:
    driver.switch_to.window(first_tab)

st.title("Hello world")

if st.button("Visit apple"):
    driver.get("https://apple.com")

if st.button("Visit google"):
    driver.get("https://google.com")

if st.button("close tab"):
    driver.close()

try:
    st.write("Current Page Title:", driver.title)
except WebDriverException:
    st.write("chrome not running")