Add colab detection util

This commit is contained in:
Cohee
2023-11-30 02:03:11 +02:00
parent aded5b4363
commit f99ea5aa89
2 changed files with 17 additions and 3 deletions

View File

@@ -4,20 +4,26 @@ from selenium.webdriver.chrome.options import Options as ChromeOptions
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.webdriver.chrome.service import Service as ChromeService
from selenium.webdriver.firefox.service import Service as FirefoxService
from modules.utils import is_colab
import atexit
def get_driver():
try:
print("Initializing Chrome driver...")
chromeService = ChromeService()
options = ChromeOptions()
options.add_argument('--disable-infobars')
options.add_argument("--headless")
options.add_argument("--disable-gpu")
options.add_argument("--no-sandbox")
options.add_argument('--disable-dev-shm-usage')
options.add_argument("--lang=en-GB")
return webdriver.Chrome(service=chromeService, options=options)
if is_colab():
return webdriver.Chrome('chromedriver', options=options)
else:
chromeService = ChromeService()
return webdriver.Chrome(service=chromeService, options=options)
except:
print("Chrome not found, using Firefox instead.")
firefoxService = FirefoxService()