correct driver typing in Session

tryolabs · Mar 1, 2024 · 2e3e60f · 2e3e60f
1 parent df3db34
commit 2e3e60f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/requestium/requestium_session.py b/requestium/requestium_session.py
@@ -3,8 +3,8 @@
 
 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 +25,7 @@ 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: str = None, headless: bool = False, default_timeout: float = 5, webdriver_options: dict = None, driver: WebDriver = None
     ) -> None:
         if webdriver_options is None:
             webdriver_options = {}