Add Driver init params (headless, chrome_user_dir) (#239)

* Add Driver init params (headless, chrome_user_dir) * fix chrome options flag syntax * small doc changes * update selenium driver package with options --------- Co-authored-by: lyie28 <[email protected]>
lavague-ai · May 18, 2024 · 24a551e · 24a551e
1 parent 0b67c8a
commit 24a551e
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 10 deletions.
diff --git a/docs/docs/get-started/quick-tour-notebook/quick-tour.ipynb b/docs/docs/get-started/quick-tour-notebook/quick-tour.ipynb
@@ -105,6 +105,7 @@
     "from lavague.core import ActionEngine\n",
     "from lavague.drivers.selenium import SeleniumDriver\n",
     "\n",
+    "# Drivers are initialized in headless by default\n",
     "selenium_driver = SeleniumDriver()\n",
     "action_engine = ActionEngine(selenium_driver)"
    ]

diff --git a/docs/docs/get-started/quick-tour.md b/docs/docs/get-started/quick-tour.md
@@ -33,7 +33,7 @@ Let's start by initializing an `ActionEngine`, which is responsible for generati
 from lavague.core import ActionEngine
 from lavague.drivers.selenium import SeleniumDriver
 
-selenium_driver = SeleniumDriver()
+selenium_driver = SeleniumDriver(headless=False)
 action_engine = ActionEngine(selenium_driver)
 ```
 

diff --git a/lavague-integrations/drivers/lavague-drivers-selenium/lavague/drivers/selenium/base.py b/lavague-integrations/drivers/lavague-drivers-selenium/lavague/drivers/selenium/base.py
@@ -17,7 +17,11 @@ def __init__(
         self,
         url: Optional[str] = None,
         get_selenium_driver: Optional[Callable[[], WebDriver]] = None,
+        headless: bool = True,
+        chrome_user_dir: Optional[str] = None,
     ):
+        self.headless = headless
+        self.chrome_user_dir = chrome_user_dir
         super().__init__(url, get_selenium_driver)
 
     def default_init_code(self) -> Any:
@@ -29,7 +33,10 @@ def default_init_code(self) -> Any:
         from selenium.webdriver.common.action_chains import ActionChains
 
         chrome_options = Options()
-        chrome_options.add_argument("--headless")
+        if self.headless:
+            chrome_options.add_argument("--headless")
+        if self.chrome_user_dir:
+            chrome_options.add_argument(f"--user-data-dir={self.chrome_user_dir}")
         chrome_options.add_argument("--no-sandbox")
 
         self.driver = webdriver.Chrome(options=chrome_options)

diff --git a/lavague-integrations/drivers/lavague-drivers-selenium/poetry.lock b/lavague-integrations/drivers/lavague-drivers-selenium/poetry.lock
diff --git a/lavague-integrations/drivers/lavague-drivers-selenium/pyproject.toml b/lavague-integrations/drivers/lavague-drivers-selenium/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "lavague-drivers-selenium"
-version = "0.1.1-post1"
+version = "0.1.2"
 description = "Selenium integration for lavague"
 authors = ["lavague-ai"]
 readme = "README.md"