Support RelativeBy in type annotations

Author: gryznar

Diff for: py/selenium/webdriver/remote/shadowroot.py

@@ -16,9 +16,15 @@
 # under the License.
 
 from hashlib import md5 as md5_hash
+from typing import List
+from typing import Optional
+from typing import Union
 
-from ..common.by import By
 from .command import Command
+from .webelement import WebElement
+from ..common.by import By
+from ..common.by import ByType
+from ..support.relative_locator import RelativeBy
 
 
 class ShadowRoot:
@@ -39,7 +45,8 @@ def __repr__(self) -> str:
             type(self), self.session.session_id, self._id
         )
 
-    def find_element(self, by: str = By.ID, value: str = None):
+    def find_element(self, by: Union[ByType, RelativeBy] = By.ID,
+                     value: Optional[str] = None) -> WebElement:
         """Find an element inside a shadow root given a By strategy and
         locator.
 
@@ -76,9 +83,11 @@ def find_element(self, by: str = By.ID, value: str = None):
             by = By.CSS_SELECTOR
             value = f'[name="{value}"]'
 
-        return self._execute(Command.FIND_ELEMENT_FROM_SHADOW_ROOT, {"using": by, "value": value})["value"]
+        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: Optional[str] = None) -> \
+    List[WebElement]:
         """Find elements inside a shadow root given a By strategy and locator.
 
         Parameters:
@@ -114,7 +123,8 @@ def find_elements(self, by: str = By.ID, value: str = None):
             by = By.CSS_SELECTOR
             value = f'[name="{value}"]'
 
-        return self._execute(Command.FIND_ELEMENTS_FROM_SHADOW_ROOT, {"using": by, "value": value})["value"]
+        return self._execute(Command.FIND_ELEMENTS_FROM_SHADOW_ROOT, {"using": by, "value": value})[
+            "value"]
 
     # Private Methods
     def _execute(self, command, params=None):

Diff for: py/selenium/webdriver/remote/webdriver.py

@@ -43,6 +43,7 @@
 from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.bidi.script import Script
 from selenium.webdriver.common.by import By
+from selenium.webdriver.common.by import ByType
 from selenium.webdriver.common.options import ArgOptions
 from selenium.webdriver.common.options import BaseOptions
 from selenium.webdriver.common.print_page_options import PrintOptions
@@ -769,7 +770,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:
@@ -805,7 +806,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:

Diff for: py/selenium/webdriver/support/event_firing_webdriver.py

@@ -17,14 +17,18 @@
 
 from typing import Any
 from typing import List
+from typing import Optional
 from typing import Tuple
+from typing import Union
 
 from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.by import By
+from selenium.webdriver.common.by import ByType
 from selenium.webdriver.remote.webdriver import WebDriver
 from selenium.webdriver.remote.webelement import WebElement
 
 from .abstract_event_listener import AbstractEventListener
+from .relative_locator import RelativeBy
 
 
 def _wrap_elements(result, ef_driver):
@@ -185,10 +189,10 @@ def clear(self) -> None:
     def send_keys(self, *value) -> None:
         self._dispatch("change_value_of", (self._webelement, self._driver), "send_keys", value)
 
-    def find_element(self, by=By.ID, value=None) -> WebElement:
+    def find_element(self, by: Union[ByType, RelativeBy] = By.ID, value: Optional[str] = None) -> WebElement:
         return self._dispatch("find", (by, value, self._driver), "find_element", (by, value))
 
-    def find_elements(self, by=By.ID, value=None) -> List[WebElement]:
+    def find_elements(self, by: Union[ByType, RelativeBy] = By.ID, value: Optional[str] = None) -> List[WebElement]:
         return self._dispatch("find", (by, value, self._driver), "find_elements", (by, value))
 
     def _dispatch(self, l_call, l_args, d_call, d_args):

Diff for: py/selenium/webdriver/support/expected_conditions.py

@@ -31,8 +31,10 @@
 from selenium.common.exceptions import StaleElementReferenceException
 from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.alert import Alert
+from selenium.webdriver.common.by import ByType
 from selenium.webdriver.remote.webdriver import WebDriver
 from selenium.webdriver.remote.webdriver import WebElement
+from selenium.webdriver.support.relative_locator import RelativeBy
 
 """
  * Canned "Expected Conditions" which are generally useful within webdriver
@@ -43,6 +45,7 @@
 T = TypeVar("T")
 
 WebDriverOrWebElement = Union[WebDriver, WebElement]
+LocatorType = Union[Tuple[ByType, str], Tuple[RelativeBy, None]]
 
 
 def title_is(title: str) -> Callable[[WebDriver], bool]:
@@ -72,7 +75,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.
 
@@ -141,7 +144,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
@@ -179,7 +182,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.
 
@@ -193,7 +196,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.
 
@@ -208,7 +211,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
@@ -231,7 +234,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.
 
@@ -249,7 +252,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.
@@ -268,7 +271,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.
@@ -288,7 +291,7 @@ def _predicate(driver: WebDriverOrWebElement):
     return _predicate
 
 
-def frame_to_be_available_and_switch_to_it(locator: Union[Tuple[str, str], str]) -> Callable[[WebDriver], bool]:
+def frame_to_be_available_and_switch_to_it(locator: Union[LocatorType, str]) -> Callable[[WebDriver], bool]:
     """An expectation for checking whether the given frame is available to
     switch to.
 
@@ -310,7 +313,7 @@ def _predicate(driver: WebDriver):
 
 
 def invisibility_of_element_located(
-    locator: Union[WebElement, Tuple[str, str]]
+    locator: Union[WebElement, LocatorType]
 ) -> Callable[[WebDriverOrWebElement], Union[WebElement, bool]]:
     """An Expectation for checking that an element is either invisible or not
     present on the DOM.
@@ -336,7 +339,7 @@ def _predicate(driver: WebDriverOrWebElement):
 
 
 def invisibility_of_element(
-    element: Union[WebElement, Tuple[str, str]]
+    element: Union[WebElement, LocatorType]
 ) -> Callable[[WebDriverOrWebElement], Union[WebElement, bool]]:
     """An Expectation for checking that an element is either invisible or not
     present on the DOM.
@@ -347,7 +350,7 @@ def invisibility_of_element(
 
 
 def element_to_be_clickable(
-    mark: Union[WebElement, Tuple[str, str]]
+    mark: Union[WebElement, LocatorType]
 ) -> Callable[[WebDriverOrWebElement], Union[Literal[False], WebElement]]:
     """An Expectation for checking an element is visible and enabled such that
     you can click it.
@@ -399,7 +402,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.
 
     locator is a tuple of (by, path)
@@ -424,7 +427,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.
@@ -474,7 +477,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.