From e0b6365846a2980169dca7cec8113b5c86fb0c72 Mon Sep 17 00:00:00 2001
From: jrycw <jerry@ycwu.space>
Date: Tue, 19 Nov 2024 04:41:21 +0800
Subject: [PATCH] Introduce `cls_driver` and `cls_wd_options` to WebDriver

---
 great_tables/_utils_selenium.py | 33 ++++++++++++---------------------
 1 file changed, 12 insertions(+), 21 deletions(-)

diff --git a/great_tables/_utils_selenium.py b/great_tables/_utils_selenium.py
index ac5a5ced1..48ff1661d 100644
--- a/great_tables/_utils_selenium.py
+++ b/great_tables/_utils_selenium.py
@@ -16,8 +16,11 @@
 
 class _BaseWebDriver:
 
-    def __init__(self):
+    def __init__(self, debug_port: bool):
+        self.debug_port = debug_port
+        self.wd_options = self.cls_wd_options()
         self.add_arguments()
+        self.driver = self.cls_driver(self.wd_options)
 
     def add_arguments(self): ...
 
@@ -34,11 +37,8 @@ def __exit__(
 
 
 class _ChromeWebDriver(_BaseWebDriver):
-    def __init__(self, debug_port: int | None = None):
-        self.debug_port = debug_port
-        self.wd_options = webdriver.ChromeOptions()
-        super().__init__()
-        self.driver = webdriver.Chrome(self.wd_options)
+    cls_driver = webdriver.Chrome
+    cls_wd_options = webdriver.ChromeOptions
 
     def add_arguments(self):
         self.wd_options.add_argument("--headless=new")
@@ -47,19 +47,13 @@ def add_arguments(self):
 
 
 class _SafariWebDriver(_BaseWebDriver):
-    def __init__(self, debug_port: int | None = None):
-        self.debug_port = debug_port
-        self.wd_options = webdriver.SafariOptions()
-        super().__init__()
-        self.driver = webdriver.Safari(self.wd_options)
+    cls_driver = webdriver.Chrome
+    cls_wd_options = webdriver.SafariOptions
 
 
 class _FirefoxWebDriver(_BaseWebDriver):
-    def __init__(self, debug_port: int | None = None):
-        self.debug_port = debug_port
-        self.wd_options = webdriver.FirefoxOptions()
-        super().__init__()
-        self.driver = webdriver.Firefox(self.wd_options)
+    cls_driver = webdriver.Firefox
+    cls_wd_options = webdriver.FirefoxOptions
 
     def add_arguments(self):
         self.wd_options.add_argument("--headless")
@@ -68,11 +62,8 @@ def add_arguments(self):
 
 
 class _EdgeWebDriver(_BaseWebDriver):
-    def __init__(self, debug_port: int | None = None):
-        self.debug_port = debug_port
-        self.wd_options = webdriver.EdgeOptions()
-        super().__init__()
-        self.driver = webdriver.Edge(self.wd_options)
+    cls_driver = webdriver.Edge
+    cls_wd_options = webdriver.EdgeOptions
 
     def add_arguments(self):
         self.wd_options.add_argument("--headless")