improved typing

requestium/requestium_chrome.py

@@ -1,7 +1,7 @@
 import functools
 import time
 import warnings
-from typing import Union
+from typing import Optional, Union
 
 import tldextract
 from parsel.selector import Selector, SelectorList
@@ -23,7 +23,7 @@ def __init__(self, *args, **kwargs) -> None:
         self.default_timeout = kwargs.pop("default_timeout", None)
         super(RequestiumChrome, self).__init__(*args, **kwargs)
 
-    def try_add_cookie(self, cookie: dict[str]) -> bool:
+    def try_add_cookie(self, cookie: dict[str, str]) -> bool:
         """
         Attempt to add the cookie. Suppress any errors, and simply
         detect success or failure if the cookie was actually added.
@@ -34,7 +34,7 @@ def try_add_cookie(self, cookie: dict[str]) -> bool:
             pass
         return self.is_cookie_in_driver(cookie)
 
-    def ensure_add_cookie(self, cookie: dict[str], override_domain: str = None) -> None:
+    def ensure_add_cookie(self, cookie: dict[str, str], override_domain: Optional[str] = None) -> None:
         """
         Ensures a cookie gets added to the driver
 
@@ -83,7 +83,7 @@ def ensure_add_cookie(self, cookie: dict[str], override_domain: str = None) -> N
             if not cookie_added:
                 raise WebDriverException(f"Couldn't add the following cookie to the webdriver: {str(cookie)}")
 
-    def is_cookie_in_driver(self, cookie: dict[str]) -> bool:
+    def is_cookie_in_driver(self, cookie: dict[str, str]) -> bool:
         """
         We check that the cookie is correctly added to the driver
 
@@ -101,31 +101,31 @@ def is_cookie_in_driver(self, cookie: dict[str]) -> bool:
                 return True
         return False
 
-    def ensure_element_by_id(self, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element_by_id(self, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         return self.ensure_element(By.ID, selector, state, timeout)
 
-    def ensure_element_by_name(self, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element_by_name(self, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         return self.ensure_element(By.NAME, selector, state, timeout)
 
-    def ensure_element_by_xpath(self, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element_by_xpath(self, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         return self.ensure_element(By.XPATH, selector, state, timeout)
 
-    def ensure_element_by_link_text(self, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element_by_link_text(self, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         return self.ensure_element(By.LINK_TEXT, selector, state, timeout)
 
-    def ensure_element_by_partial_link_text(self, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element_by_partial_link_text(self, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         return self.ensure_element(By.PARTIAL_LINK_TEXT, selector, state, timeout)
 
-    def ensure_element_by_tag_name(self, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element_by_tag_name(self, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         return self.ensure_element(By.TAG_NAME, selector, state, timeout)
 
-    def ensure_element_by_class_name(self, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element_by_class_name(self, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         return self.ensure_element(By.CLASS_NAME, selector, state, timeout)
 
-    def ensure_element_by_css_selector(self, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element_by_css_selector(self, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         return self.ensure_element(By.CSS_SELECTOR, selector, state, timeout)
 
-    def ensure_element(self, locator: str, selector: str, state: str = "present", timeout: float = None) -> WebElement:
+    def ensure_element(self, locator: str, selector: str, state: str = "present", timeout: Union[float | int] = 1) -> Optional[WebElement]:
         """
         This method allows us to wait till an element appears or disappears in the browser
 

requestium/requestium_session.py

@@ -1,10 +1,12 @@
 import functools
 import types
+from typing import Optional, Union
 
 import requests
 import tldextract
 from selenium import webdriver
 from selenium.webdriver import ChromeService
+from selenium.webdriver.remote.webdriver import WebDriver
 
 from .requestium_chrome import RequestiumChrome
 from .requestium_response import RequestiumResponse
@@ -25,7 +27,12 @@ class Session(requests.Session):
     """
 
     def __init__(
-        self, webdriver_path: str = None, headless: bool = False, default_timeout: float = 5, webdriver_options: dict = None, driver: webdriver = None
+        self,
+        webdriver_path: Optional[str] = None,
+        headless: bool = False,
+        default_timeout: Union[float | int] = 5,
+        webdriver_options: Optional[dict] = None,
+        driver: Optional[WebDriver] = None,
     ) -> None:
         if webdriver_options is None:
             webdriver_options = {}
@@ -68,21 +75,24 @@ def _start_chrome_browser(self, headless: bool = False) -> RequestiumChrome:
 
         args = self.webdriver_options.get("arguments")
         if isinstance(args, list):
-            [chrome_options.add_argument(arg) for arg in args]
+            for arg in args:
+                chrome_options.add_argument(arg)
         elif args:
             raise Exception(f"A list is needed to use 'arguments' option. Found {type(args)}")
 
         extensions = self.webdriver_options.get("extensions")
         if isinstance(extensions, list):
-            [chrome_options.add_extension(arg) for arg in extensions]
+            for arg in extensions:
+                chrome_options.add_extension(arg)
 
         if "prefs" in self.webdriver_options:
             prefs = self.webdriver_options["prefs"]
             chrome_options.add_experimental_option("prefs", prefs)
 
         experimental_options = self.webdriver_options.get("experimental_options")
         if isinstance(experimental_options, dict):
-            [chrome_options.add_experimental_option(name, value) for name, value in experimental_options.items()]
+            for name, value in experimental_options.items():
+                chrome_options.add_experimental_option(name, value)
 
         # Selenium updated webdriver.Chrome's arg and kwargs, to accept options, service, keep_alive
         # since ChromeService is the only object where webdriver_path is mapped to executable_path, it must be
@@ -95,7 +105,7 @@ def _start_chrome_browser(self, headless: bool = False) -> RequestiumChrome:
             default_timeout=self.default_timeout,
         )
 
-    def transfer_session_cookies_to_driver(self, domain: str = None) -> None:
+    def transfer_session_cookies_to_driver(self, domain: Optional[str] = None) -> None:
         """
         Copies the Session's cookies into the webdriver