Fix: Add support for RelativeBy in type annotations

Author: ShauryaDusht

py/selenium/webdriver/remote/shadowroot.py

@@ -16,8 +16,11 @@
 # under the License.
 
 from hashlib import md5 as md5_hash
+from typing import Union, TYPE_CHECKING
+if TYPE_CHECKING:
+    from selenium.webdriver.support.relative_locator import RelativeBy
 
-from ..common.by import By
+from ..common.by import By, ByType
 from .command import Command
 
 
@@ -43,7 +46,7 @@ def __repr__(self) -> str:
     def id(self) -> str:
         return self._id
 
-    def find_element(self, by: str = By.ID, value: str = None):
+    def find_element(self, by: "Union[ByType, RelativeBy]" = By.ID, value: str = None):
         """Find an element inside a shadow root given a By strategy and
         locator.
 
@@ -82,7 +85,7 @@ def find_element(self, by: str = By.ID, value: str = None):
 
         return self._execute(Command.FIND_ELEMENT_FROM_SHADOW_ROOT, {"using": by, "value": value})["value"]
 
-    def find_elements(self, by: str = By.ID, value: str = None):
+    def find_elements(self, by: "Union[ByType, RelativeBy]" = By.ID, value: str = None):
         """Find elements inside a shadow root given a By strategy and locator.
 
         Parameters:

py/selenium/webdriver/remote/webdriver.py

@@ -30,7 +30,10 @@
 from base64 import b64decode, urlsafe_b64encode
 from contextlib import asynccontextmanager, contextmanager
 from importlib import import_module
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from selenium.webdriver.support.relative_locator import RelativeBy
 
 from selenium.common.exceptions import (
     InvalidArgumentException,
@@ -46,7 +49,7 @@
 from selenium.webdriver.common.bidi.session import Session
 from selenium.webdriver.common.bidi.storage import Storage
 from selenium.webdriver.common.bidi.webextension import WebExtension
-from selenium.webdriver.common.by import By
+from selenium.webdriver.common.by import By, ByType
 from selenium.webdriver.common.options import ArgOptions, BaseOptions
 from selenium.webdriver.common.print_page_options import PrintOptions
 from selenium.webdriver.common.timeouts import Timeouts
@@ -55,8 +58,6 @@
     VirtualAuthenticatorOptions,
     required_virtual_authenticator,
 )
-from selenium.webdriver.support.relative_locator import RelativeBy
-
 from ..common.fedcm.dialog import Dialog
 from .bidi_connection import BidiConnection
 from .client_config import ClientConfig
@@ -877,7 +878,7 @@ def timeouts(self, timeouts) -> None:
         """
         _ = self.execute(Command.SET_TIMEOUTS, timeouts._to_json())["value"]
 
-    def find_element(self, by=By.ID, value: Optional[str] = None) -> WebElement:
+    def find_element(self, by: "Union[ByType, RelativeBy]" = By.ID, value: Optional[str] = None) -> WebElement:
         """Find an element given a By strategy and locator.
 
         Parameters:
@@ -913,7 +914,7 @@ def find_element(self, by=By.ID, value: Optional[str] = None) -> WebElement:
 
         return self.execute(Command.FIND_ELEMENT, {"using": by, "value": value})["value"]
 
-    def find_elements(self, by=By.ID, value: Optional[str] = None) -> list[WebElement]:
+    def find_elements(self, by: "Union[ByType, RelativeBy]" = By.ID, value: Optional[str] = None) -> list[WebElement]:
         """Find elements given a By strategy and locator.
 
         Parameters:

py/selenium/webdriver/remote/webelement.py

@@ -24,9 +24,14 @@
 from base64 import b64decode, encodebytes
 from hashlib import md5 as md5_hash
 from io import BytesIO
+from typing import Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from selenium.webdriver.support.relative_locator import RelativeBy
+
 
 from selenium.common.exceptions import JavascriptException, WebDriverException
-from selenium.webdriver.common.by import By
+from selenium.webdriver.common.by import By, ByType
 from selenium.webdriver.common.utils import keys_to_typing
 
 from .command import Command
@@ -572,7 +577,7 @@ def _execute(self, command, params=None):
         params["id"] = self._id
         return self._parent.execute(command, params)
 
-    def find_element(self, by=By.ID, value=None) -> WebElement:
+    def find_element(self, by: "Union[ByType, RelativeBy]" = By.ID, value=None) -> WebElement:
         """Find an element given a By strategy and locator.
 
         Parameters:
@@ -601,7 +606,7 @@ def find_element(self, by=By.ID, value=None) -> WebElement:
         by, value = self._parent.locator_converter.convert(by, value)
         return self._execute(Command.FIND_CHILD_ELEMENT, {"using": by, "value": value})["value"]
 
-    def find_elements(self, by=By.ID, value=None) -> list[WebElement]:
+    def find_elements(self, by: "Union[ByType, RelativeBy]" = By.ID, value=None) -> list[WebElement]:
         """Find elements given a By strategy and locator.
 
         Parameters:

py/selenium/webdriver/support/expected_conditions.py

@@ -17,7 +17,10 @@
 
 import re
 from collections.abc import Iterable
-from typing import Any, Callable, Literal, TypeVar, Union
+from typing import Any, Callable, Literal, TypeVar, Union, Tuple
+
+from selenium.webdriver.common.by import ByType
+from selenium.webdriver.support.relative_locator import RelativeBy
 
 from selenium.common.exceptions import (
     NoAlertPresentException,
@@ -38,6 +41,7 @@
 T = TypeVar("T")
 
 WebDriverOrWebElement = Union[WebDriver, WebElement]
+LocatorType = Union[Tuple[ByType, str], Tuple[RelativeBy, None]]
 
 
 def title_is(title: str) -> Callable[[WebDriver], bool]:
@@ -79,7 +83,7 @@ def _predicate(driver: WebDriver):
     return _predicate
 
 
-def presence_of_element_located(locator: tuple[str, str]) -> Callable[[WebDriverOrWebElement], WebElement]:
+def presence_of_element_located(locator: LocatorType) -> Callable[[WebDriverOrWebElement], WebElement]:
     """An expectation for checking that an element is present on the DOM of a
     page. This does not necessarily mean that the element is visible.
 
@@ -189,7 +193,7 @@ def _predicate(driver: WebDriver):
 
 
 def visibility_of_element_located(
-    locator: tuple[str, str],
+    locator: LocatorType,
 ) -> Callable[[WebDriverOrWebElement], Union[Literal[False], WebElement]]:
     """An expectation for checking that an element is present on the DOM of a
     page and visible. Visibility means that the element is not only displayed
@@ -272,7 +276,7 @@ def _element_if_visible(element: WebElement, visibility: bool = True) -> Union[L
     return element if element.is_displayed() == visibility else False
 
 
-def presence_of_all_elements_located(locator: tuple[str, str]) -> Callable[[WebDriverOrWebElement], list[WebElement]]:
+def presence_of_all_elements_located(locator: LocatorType) -> Callable[[WebDriverOrWebElement], list[WebElement]]:
     """An expectation for checking that there is at least one element present
     on a web page.
 
@@ -299,7 +303,7 @@ def _predicate(driver: WebDriverOrWebElement):
     return _predicate
 
 
-def visibility_of_any_elements_located(locator: tuple[str, str]) -> Callable[[WebDriverOrWebElement], list[WebElement]]:
+def visibility_of_any_elements_located(locator: LocatorType) -> Callable[[WebDriverOrWebElement], list[WebElement]]:
     """An expectation for checking that there is at least one element visible
     on a web page.
 
@@ -327,7 +331,7 @@ def _predicate(driver: WebDriverOrWebElement):
 
 
 def visibility_of_all_elements_located(
-    locator: tuple[str, str],
+    locator: LocatorType,
 ) -> Callable[[WebDriverOrWebElement], Union[list[WebElement], Literal[False]]]:
     """An expectation for checking that all elements are present on the DOM of
     a page and visible. Visibility means that the elements are not only
@@ -363,7 +367,7 @@ def _predicate(driver: WebDriverOrWebElement):
     return _predicate
 
 
-def text_to_be_present_in_element(locator: tuple[str, str], text_: str) -> Callable[[WebDriverOrWebElement], bool]:
+def text_to_be_present_in_element(locator: LocatorType, text_: str) -> Callable[[WebDriverOrWebElement], bool]:
     """An expectation for checking if the given text is present in the
     specified element.
 
@@ -399,7 +403,7 @@ def _predicate(driver: WebDriverOrWebElement):
 
 
 def text_to_be_present_in_element_value(
-    locator: tuple[str, str], text_: str
+    locator: LocatorType, text_: str
 ) -> Callable[[WebDriverOrWebElement], bool]:
     """An expectation for checking if the given text is present in the
     element's value.
@@ -436,7 +440,7 @@ def _predicate(driver: WebDriverOrWebElement):
 
 
 def text_to_be_present_in_element_attribute(
-    locator: tuple[str, str], attribute_: str, text_: str
+    locator: LocatorType, attribute_: str, text_: str
 ) -> Callable[[WebDriverOrWebElement], bool]:
     """An expectation for checking if the given text is present in the
     element's attribute.
@@ -687,7 +691,7 @@ def _predicate(_):
     return _predicate
 
 
-def element_located_to_be_selected(locator: tuple[str, str]) -> Callable[[WebDriverOrWebElement], bool]:
+def element_located_to_be_selected(locator: LocatorType) -> Callable[[WebDriverOrWebElement], bool]:
     """An expectation for the element to be located is selected.
 
     Parameters:
@@ -743,7 +747,7 @@ def _predicate(_):
 
 
 def element_located_selection_state_to_be(
-    locator: tuple[str, str], is_selected: bool
+    locator: LocatorType, is_selected: bool
 ) -> Callable[[WebDriverOrWebElement], bool]:
     """An expectation to locate an element and check if the selection state
     specified is in that state.
@@ -858,7 +862,7 @@ def _predicate(driver: WebDriver):
     return _predicate
 
 
-def element_attribute_to_include(locator: tuple[str, str], attribute_: str) -> Callable[[WebDriverOrWebElement], bool]:
+def element_attribute_to_include(locator: LocatorType, attribute_: str) -> Callable[[WebDriverOrWebElement], bool]:
     """An expectation for checking if the given attribute is included in the
     specified element.
 

py/test/unit/selenium/webdriver/support/test_expected_conditions_relative_by.py

@@ -0,0 +1,57 @@
+import pytest
+from selenium.webdriver.support import expected_conditions as EC
+from selenium.webdriver.support.relative_locator import locate_with
+from selenium.webdriver.common.by import By
+from unittest.mock import Mock
+
+
+class TestExpectedConditionsRelativeBy:
+    """Test that expected conditions accept RelativeBy in type annotations"""
+    
+    def test_presence_of_element_located_accepts_relative_by(self):
+        """Test presence_of_element_located accepts RelativeBy"""
+        relative_by = locate_with(By.TAG_NAME, "div").above({By.ID: "footer"})
+        condition = EC.presence_of_element_located(relative_by)
+        assert condition is not None
+    
+    def test_visibility_of_element_located_accepts_relative_by(self):
+        """Test visibility_of_element_located accepts RelativeBy"""
+        relative_by = locate_with(By.TAG_NAME, "button").near({By.CLASS_NAME: "submit"})
+        condition = EC.visibility_of_element_located(relative_by)
+        assert condition is not None
+    
+    def test_presence_of_all_elements_located_accepts_relative_by(self):
+        """Test presence_of_all_elements_located accepts RelativeBy"""
+        relative_by = locate_with(By.CSS_SELECTOR, ".item").below({By.ID: "header"})
+        condition = EC.presence_of_all_elements_located(relative_by)
+        assert condition is not None
+    
+    def test_visibility_of_any_elements_located_accepts_relative_by(self):
+        """Test visibility_of_any_elements_located accepts RelativeBy"""
+        relative_by = locate_with(By.TAG_NAME, "span").to_left_of({By.ID: "sidebar"})
+        condition = EC.visibility_of_any_elements_located(relative_by)
+        assert condition is not None
+    
+    def test_text_to_be_present_in_element_accepts_relative_by(self):
+        """Test text_to_be_present_in_element accepts RelativeBy"""
+        relative_by = locate_with(By.TAG_NAME, "p").above({By.CLASS_NAME: "footer"})
+        condition = EC.text_to_be_present_in_element(relative_by, "Hello")
+        assert condition is not None
+    
+    def test_element_to_be_clickable_accepts_relative_by(self):
+        """Test element_to_be_clickable accepts RelativeBy"""
+        relative_by = locate_with(By.TAG_NAME, "button").near({By.ID: "form"})
+        condition = EC.element_to_be_clickable(relative_by)
+        assert condition is not None
+    
+    def test_invisibility_of_element_located_accepts_relative_by(self):
+        """Test invisibility_of_element_located accepts RelativeBy"""
+        relative_by = locate_with(By.CSS_SELECTOR, ".loading").above({By.ID: "content"})
+        condition = EC.invisibility_of_element_located(relative_by)
+        assert condition is not None
+    
+    def test_element_located_to_be_selected_accepts_relative_by(self):
+        """Test element_located_to_be_selected accepts RelativeBy"""
+        relative_by = locate_with(By.TAG_NAME, "input").near({By.ID: "terms-label"})
+        condition = EC.element_located_to_be_selected(relative_by)
+        assert condition is not None

py/test/unit/selenium/webdriver/test_relative_by_annotations.py

@@ -0,0 +1,65 @@
+import pytest
+from selenium.webdriver.common.by import By
+from selenium.webdriver.support.relative_locator import RelativeBy, locate_with
+from selenium.webdriver.remote.webdriver import WebDriver
+from selenium.webdriver.remote.webelement import WebElement
+from selenium.webdriver.remote.shadowroot import ShadowRoot
+from unittest.mock import Mock, MagicMock
+
+
+class TestRelativeByAnnotations:
+    """Test that RelativeBy is properly accepted in type annotations"""
+    
+    def test_webdriver_find_element_accepts_relative_by(self):
+        """Test WebDriver.find_element accepts RelativeBy"""
+        driver = Mock(spec=WebDriver)
+        relative_by = locate_with(By.TAG_NAME, "div").above({By.ID: "footer"})
+        
+        # This should not raise type checking errors
+        driver.find_element(by=relative_by)
+        driver.find_element(relative_by)
+    
+    def test_webdriver_find_elements_accepts_relative_by(self):
+        """Test WebDriver.find_elements accepts RelativeBy"""
+        driver = Mock(spec=WebDriver)
+        relative_by = locate_with(By.TAG_NAME, "div").below({By.ID: "header"})
+        
+        # This should not raise type checking errors
+        driver.find_elements(by=relative_by)
+        driver.find_elements(relative_by)
+    
+    def test_webelement_find_element_accepts_relative_by(self):
+        """Test WebElement.find_element accepts RelativeBy"""
+        element = Mock(spec=WebElement)
+        relative_by = locate_with(By.TAG_NAME, "span").near({By.CLASS_NAME: "button"})
+        
+        # This should not raise type checking errors
+        element.find_element(by=relative_by)
+        element.find_element(relative_by)
+    
+    def test_webelement_find_elements_accepts_relative_by(self):
+        """Test WebElement.find_elements accepts RelativeBy"""
+        element = Mock(spec=WebElement)
+        relative_by = locate_with(By.TAG_NAME, "input").to_left_of({By.ID: "submit"})
+        
+        # This should not raise type checking errors
+        element.find_elements(by=relative_by)
+        element.find_elements(relative_by)
+    
+    def test_shadowroot_find_element_accepts_relative_by(self):
+        """Test ShadowRoot.find_element accepts RelativeBy"""
+        shadow_root = Mock(spec=ShadowRoot)
+        relative_by = locate_with(By.TAG_NAME, "button").to_right_of({By.ID: "cancel"})
+        
+        # This should not raise type checking errors
+        shadow_root.find_element(by=relative_by)
+        shadow_root.find_element(relative_by)
+    
+    def test_shadowroot_find_elements_accepts_relative_by(self):
+        """Test ShadowRoot.find_elements accepts RelativeBy"""
+        shadow_root = Mock(spec=ShadowRoot)
+        relative_by = locate_with(By.CSS_SELECTOR, ".item").above({By.ID: "footer"})
+        
+        # This should not raise type checking errors
+        shadow_root.find_elements(by=relative_by)
+        shadow_root.find_elements(relative_by)