Skip to content

Commit

Permalink
improved typing
Browse files Browse the repository at this point in the history
  • Loading branch information
bmos committed Mar 1, 2024
1 parent 9561710 commit e00e5cd
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 18 deletions.
26 changes: 13 additions & 13 deletions requestium/requestium_chrome.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
20 changes: 15 additions & 5 deletions requestium/requestium_session.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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 = {}
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down

0 comments on commit e00e5cd

Please sign in to comment.