Add IntelliSense autocompletion support

Author: mdmintz

seleniumbase/core/sb_driver.py

@@ -1,5 +1,6 @@
 """Add new methods to extend the driver"""
 from contextlib import suppress
+from selenium.webdriver.remote.webdriver import WebDriver
 from selenium.webdriver.remote.webelement import WebElement
 from seleniumbase.config import settings
 from seleniumbase.fixtures import js_utils
@@ -8,7 +9,7 @@
 from seleniumbase.fixtures import shared_utils
 
 
-class DriverMethods():
+class DriverMethods(WebDriver):
     def __init__(self, driver):
         self.driver = driver
 

seleniumbase/plugins/driver_manager.py

@@ -38,6 +38,7 @@
 """
 import os
 import sys
+from seleniumbase.core import sb_driver
 
 
 class DriverContext():
@@ -139,7 +140,7 @@ def Driver(
     pls=None,  # Shortcut / Duplicate of "page_load_strategy".
     cft=None,  # Use "Chrome for Testing"
     chs=None,  # Use "Chrome-Headless-Shell"
-):
+) -> sb_driver.DriverMethods:
     """
     * SeleniumBase Driver as a Python Context Manager or a returnable object. *
 

seleniumbase/plugins/sb_manager.py

@@ -24,6 +24,8 @@
 #########################################
 """
 from contextlib import contextmanager, suppress
+from typing import Any, Generator
+from seleniumbase import BaseCase
 
 
 @contextmanager  # Usage: -> ``with SB() as sb:``
@@ -133,7 +135,7 @@ def SB(
     highlights=None,  # Number of highlight animations for Demo Mode actions.
     interval=None,  # SECONDS (Autoplay interval for SB Slides & Tour steps.)
     time_limit=None,  # SECONDS (Safely fail tests that exceed the time limit.)
-):
+) -> Generator[BaseCase, Any, None]:
     """
     * SeleniumBase as a Python Context Manager *
 
@@ -263,7 +265,6 @@ def SB(
     import sys
     import time
     import traceback
-    from seleniumbase import BaseCase
     from seleniumbase import config as sb_config
     from seleniumbase.config import settings
     from seleniumbase.fixtures import constants