summaryrefslogtreecommitdiffstats
path: root/g4f/Provider/webdriver.py
diff options
context:
space:
mode:
authorTekky <98614666+xtekky@users.noreply.github.com>2023-11-20 19:27:38 +0100
committerGitHub <noreply@github.com>2023-11-20 19:27:38 +0100
commite8d88c955f75f539dd71bd4b713e90094751161c (patch)
tree6d1ad2636abfd7ad0b4f5a59aa4630cf2a29723e /g4f/Provider/webdriver.py
parentMerge pull request #1275 from egcash/patch-1 (diff)
parentMerge branch 'main' into webdriver (diff)
downloadgpt4free-e8d88c955f75f539dd71bd4b713e90094751161c.tar
gpt4free-e8d88c955f75f539dd71bd4b713e90094751161c.tar.gz
gpt4free-e8d88c955f75f539dd71bd4b713e90094751161c.tar.bz2
gpt4free-e8d88c955f75f539dd71bd4b713e90094751161c.tar.lz
gpt4free-e8d88c955f75f539dd71bd4b713e90094751161c.tar.xz
gpt4free-e8d88c955f75f539dd71bd4b713e90094751161c.tar.zst
gpt4free-e8d88c955f75f539dd71bd4b713e90094751161c.zip
Diffstat (limited to '')
-rw-r--r--g4f/Provider/webdriver.py92
1 files changed, 92 insertions, 0 deletions
diff --git a/g4f/Provider/webdriver.py b/g4f/Provider/webdriver.py
new file mode 100644
index 00000000..da3b13ed
--- /dev/null
+++ b/g4f/Provider/webdriver.py
@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import time
+from platformdirs import user_config_dir
+try:
+ from selenium.webdriver.remote.webdriver import WebDriver
+except ImportError:
+ class WebDriver():
+ pass
+try:
+ from undetected_chromedriver import Chrome, ChromeOptions
+except ImportError:
+ class Chrome():
+ def __init__():
+ raise RuntimeError('Please install the "undetected_chromedriver" package')
+ class ChromeOptions():
+ def add_argument():
+ pass
+try:
+ from pyvirtualdisplay import Display
+ has_pyvirtualdisplay = True
+except ImportError:
+ has_pyvirtualdisplay = False
+
+def get_browser(
+ user_data_dir: str = None,
+ headless: bool = False,
+ proxy: str = None,
+ options: ChromeOptions = None
+) -> Chrome:
+ if user_data_dir == None:
+ user_data_dir = user_config_dir("g4f")
+ if proxy:
+ if not options:
+ options = ChromeOptions()
+ options.add_argument(f'--proxy-server={proxy}')
+ return Chrome(options=options, user_data_dir=user_data_dir, headless=headless)
+
+class WebDriverSession():
+ def __init__(
+ self,
+ webdriver: WebDriver = None,
+ user_data_dir: str = None,
+ headless: bool = False,
+ virtual_display: bool = False,
+ proxy: str = None,
+ options: ChromeOptions = None
+ ):
+ self.webdriver = webdriver
+ self.user_data_dir = user_data_dir
+ self.headless = headless
+ self.virtual_display = None
+ if has_pyvirtualdisplay and virtual_display:
+ self.virtual_display = Display(size=(1920,1080))
+ self.proxy = proxy
+ self.options = options
+ self.default_driver = None
+
+ def reopen(
+ self,
+ user_data_dir: str = None,
+ headless: bool = False,
+ virtual_display: bool = False
+ ) -> WebDriver:
+ if user_data_dir == None:
+ user_data_dir = self.user_data_dir
+ if self.default_driver:
+ self.default_driver.quit()
+ if not virtual_display and self.virtual_display:
+ self.virtual_display.stop()
+ self.virtual_display = None
+ self.default_driver = get_browser(user_data_dir, headless, self.proxy)
+ return self.default_driver
+
+ def __enter__(self) -> WebDriver:
+ if self.webdriver:
+ return self.webdriver
+ if self.virtual_display:
+ self.virtual_display.start()
+ self.default_driver = get_browser(self.user_data_dir, self.headless, self.proxy, self.options)
+ return self.default_driver
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ if self.default_driver:
+ try:
+ self.default_driver.close()
+ except:
+ pass
+ time.sleep(0.1)
+ self.default_driver.quit()
+ if self.virtual_display:
+ self.virtual_display.stop() \ No newline at end of file