Added FireFox Support

uNickz · Dec 23, 2023 · 1211ea0 · 1211ea0
1 parent 279414c
commit 1211ea0
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 18 deletions.
diff --git a/WebSite2PDF/client.py b/WebSite2PDF/client.py
@@ -3,8 +3,9 @@
 import os
 import re
 
-from selenium.webdriver.chrome.options import Options
-from typing import Optional, Iterable, Union, List
+from selenium.webdriver.firefox.options import Options as FireFoxOptions
+from selenium.webdriver.chrome.options import Options as ChromeOptions
+from typing import Optional, Iterable, Union, List, Tuple
 from .driver import Driver
 
 log = logging.getLogger(__name__)
@@ -44,7 +45,7 @@ def pdf(self, url: Union[str, Iterable[str]], filename: Optional[Union[str, Iter
         __seleniumOptions: List[str] = seleniumOptions if seleniumOptions else self.seleniumOptions
         __delay: int = delay if delay else self.delay
 
-        __opts: Options = self.__FormatOptions(__seleniumOptions)
+        __opts: Tuple[ChromeOptions, FireFoxOptions] = self.__FormatOptions(__seleniumOptions)
         self.client.start_client(__opts)
         if isinstance(url, str):
             self.client.connect(url)
@@ -91,9 +92,11 @@ def pdf(self, url: Union[str, Iterable[str]], filename: Optional[Union[str, Iter
         self.client.stop_client()
         return __gresp
 
-    def __FormatOptions(self, seleniumOptions: List[str]) -> Options:
-        __opts = Options()
+    def __FormatOptions(self, seleniumOptions: List[str]) -> Tuple[ChromeOptions, FireFoxOptions]:
+        __chrome_opts = ChromeOptions()
+        __firefox_opts = FireFoxOptions()
         for _ in seleniumOptions:
-            __opts.add_argument(_)
-        __opts.add_experimental_option("prefs", {"profile.managed_default_content_settings.images" : 2})
-        return __opts
+            __chrome_opts.add_argument(_)
+            __firefox_opts.add_argument(_)
+        __chrome_opts.add_experimental_option("prefs", {"profile.managed_default_content_settings.images" : 2})
+        return (__chrome_opts, __firefox_opts)
diff --git a/WebSite2PDF/driver/driver.py b/WebSite2PDF/driver/driver.py
@@ -4,16 +4,18 @@
 import json
 import os
 
+from selenium.webdriver.firefox.options import Options as FireFoxOptions
+from selenium.webdriver.firefox.service import Service as FireFoxService
 from selenium.webdriver.support.expected_conditions import staleness_of
+from selenium.webdriver.chrome.options import Options as ChromeOptions
+from selenium.webdriver.chrome.service import Service as ChromeService
 from webdriver_manager.chrome import ChromeDriverManager
 from webdriver_manager.firefox import GeckoDriverManager
 from selenium.common.exceptions import TimeoutException
 from selenium.webdriver.support.ui import WebDriverWait
-from selenium.webdriver.chrome.service import Service
-from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.by import By
 from selenium import webdriver
-from typing import Optional
+from typing import Optional, Tuple
 
 from ..errors import ClientAlreadyStarted, ClientAlreadyStopped, InvalidUrl, RequestFailed, InvalidFile
 
@@ -23,26 +25,26 @@ class Driver:
     def __init__(self) -> None:
         self.driver: webdriver.Chrome = None
 
-    def start_client(self, options: Optional[Options] = None, service: Optional[Service] = None, install_driver: bool = True) -> None:
+    def start_client(self, options: Optional[Tuple[ChromeOptions, FireFoxOptions]] = None, service: Optional[Tuple[FireFoxService, ChromeService]] = None, install_driver: bool = True) -> None:
         if self.driver:
             raise ClientAlreadyStarted
 
         if install_driver:
             if tuple(map(int, selenium.__version__.split("."))) >= (4, 10, 0):
                 try:
-                    self.driver = webdriver.Chrome(options = options, service = Service(executable_path = ChromeDriverManager().install()))
+                    self.driver = webdriver.Chrome(options = options[0] if options else None, service = ChromeService(executable_path = ChromeDriverManager().install()))
                 except AttributeError:
-                    self.driver = webdriver.Firefox(options = options, service = Service(executable_path = GeckoDriverManager().install()))
+                    self.driver = webdriver.Firefox(options = options[1] if options else None, service = FireFoxService(executable_path = GeckoDriverManager().install()))
             else:
                 try:
-                    self.driver = webdriver.Chrome(executable_path = ChromeDriverManager().install(), options = options, service = service)
+                    self.driver = webdriver.Chrome(executable_path = ChromeDriverManager().install(), options = options[0] if options else None, service = service[0] if service else None)
                 except AttributeError:
-                    self.driver = webdriver.Firefox(executable_path = GeckoDriverManager().install(), options = options, service = service)
+                    self.driver = webdriver.Firefox(executable_path = GeckoDriverManager().install(), options = options[1] if options else None, service = service[1] if service else None)
         else:
             try:
-                self.driver = webdriver.Chrome(options = options, service = service)
+                self.driver = webdriver.Chrome(options = options[0] if options else None, service = service[0] if service else None)
             except AttributeError:
-                self.driver = webdriver.Firefox(options = options, service = service)
+                self.driver = webdriver.Firefox(options = options[1] if options else None, service = service[1] if service else None)
 
     def stop_client(self) -> None:
         if not self.driver: