Skip to content

Commit

Permalink
py: fix WebDriverWait type hints for WebElement (#13108)
Browse files Browse the repository at this point in the history 
* fix type hint of WebDriverWait to accept WebElement

* fix some type hint of expected_conditions to accept WebElement

* linting

---------

Co-authored-by: Diego Molina <[email protected]>
  • Loading branch information
@pinterior @diemol
pinterior and diemol authored Nov 9, 2023
1 parent 1310bb6 commit 10adfe8
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 44 deletions.
88 changes: 49 additions & 39 deletions py/selenium/webdriver/support/expected_conditions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,12 +34,16 @@
from selenium.webdriver.remote.webdriver import WebDriver
from selenium.webdriver.remote.webdriver import WebElement

T = TypeVar("T")
"""
* Canned "Expected Conditions" which are generally useful within webdriver
* tests.
"""

D = TypeVar("D")
T = TypeVar("T")

WebDriverOrWebElement = Union[WebDriver, WebElement]


def title_is(title: str) -> Callable[[WebDriver], bool]:
"""An expectation for checking the title of a page.
Expand Down Expand Up @@ -68,15 +72,15 @@ def _predicate(driver: WebDriver):
return _predicate


def presence_of_element_located(locator: Tuple[str, str]) -> Callable[[WebDriver], WebElement]:
def presence_of_element_located(locator: Tuple[str, str]) -> 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.
locator - used to find the element
returns the WebElement once it is located
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
return driver.find_element(*locator)

return _predicate
Expand Down Expand Up @@ -136,7 +140,9 @@ def _predicate(driver: WebDriver):
return _predicate


def visibility_of_element_located(locator: Tuple[str, str]) -> Callable[[WebDriver], Union[Literal[False], WebElement]]:
def visibility_of_element_located(
locator: Tuple[str, str]
) -> 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
but also has a height and width that is greater than 0.
Expand All @@ -145,7 +151,7 @@ def visibility_of_element_located(locator: Tuple[str, str]) -> Callable[[WebDriv
returns the WebElement once it is located and visible
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
try:
return _element_if_visible(driver.find_element(*locator))
except StaleElementReferenceException:
Expand All @@ -154,7 +160,7 @@ def _predicate(driver: WebDriver):
return _predicate


def visibility_of(element: WebElement) -> Callable[[WebDriver], Union[Literal[False], WebElement]]:
def visibility_of(element: WebElement) -> Callable[[Any], Union[Literal[False], WebElement]]:
"""An expectation for checking that an element, known to be present on the
DOM of a page, is visible.
Expand All @@ -173,37 +179,37 @@ 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[[WebDriver], List[WebElement]]:
def presence_of_all_elements_located(locator: Tuple[str, str]) -> Callable[[WebDriverOrWebElement], List[WebElement]]:
"""An expectation for checking that there is at least one element present
on a web page.
locator is used to find the element returns the list of WebElements
once they are located
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
return driver.find_elements(*locator)

return _predicate


def visibility_of_any_elements_located(locator: Tuple[str, str]) -> Callable[[WebDriver], List[WebElement]]:
def visibility_of_any_elements_located(locator: Tuple[str, str]) -> Callable[[WebDriverOrWebElement], List[WebElement]]:
"""An expectation for checking that there is at least one element visible
on a web page.
locator is used to find the element returns the list of WebElements
once they are located
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
return [element for element in driver.find_elements(*locator) if _element_if_visible(element)]

return _predicate


def visibility_of_all_elements_located(
locator: Tuple[str, str]
) -> Callable[[WebDriver], Union[List[WebElement], Literal[False]]]:
) -> 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
displayed but also has a height and width that is greater than 0.
Expand All @@ -212,7 +218,7 @@ def visibility_of_all_elements_located(
returns the list of WebElements once they are located and visible
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
try:
elements = driver.find_elements(*locator)
for element in elements:
Expand All @@ -225,14 +231,14 @@ def _predicate(driver: WebDriver):
return _predicate


def text_to_be_present_in_element(locator: Tuple[str, str], text_: str) -> Callable[[WebDriver], bool]:
def text_to_be_present_in_element(locator: Tuple[str, str], text_: str) -> Callable[[WebDriverOrWebElement], bool]:
"""An expectation for checking if the given text is present in the
specified element.
locator, text
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
try:
element_text = driver.find_element(*locator).text
return text_ in element_text
Expand All @@ -242,14 +248,16 @@ def _predicate(driver: WebDriver):
return _predicate


def text_to_be_present_in_element_value(locator: Tuple[str, str], text_: str) -> Callable[[WebDriver], bool]:
def text_to_be_present_in_element_value(
locator: Tuple[str, str], text_: str
) -> Callable[[WebDriverOrWebElement], bool]:
"""An expectation for checking if the given text is present in the
element's value.
locator, text
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
try:
element_text = driver.find_element(*locator).get_attribute("value")
return text_ in element_text
Expand All @@ -261,14 +269,14 @@ def _predicate(driver: WebDriver):

def text_to_be_present_in_element_attribute(
locator: Tuple[str, str], attribute_: str, text_: str
) -> Callable[[WebDriver], bool]:
) -> Callable[[WebDriverOrWebElement], bool]:
"""An expectation for checking if the given text is present in the
element's attribute.
locator, attribute, text
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
try:
element_text = driver.find_element(*locator).get_attribute(attribute_)
if element_text is None:
Expand Down Expand Up @@ -303,14 +311,14 @@ def _predicate(driver: WebDriver):

def invisibility_of_element_located(
locator: Union[WebElement, Tuple[str, str]]
) -> Callable[[WebDriver], Union[WebElement, bool]]:
) -> Callable[[WebDriverOrWebElement], Union[WebElement, bool]]:
"""An Expectation for checking that an element is either invisible or not
present on the DOM.
locator used to find the element
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
try:
target = locator
if not isinstance(target, WebElement):
Expand All @@ -329,7 +337,7 @@ def _predicate(driver: WebDriver):

def invisibility_of_element(
element: Union[WebElement, Tuple[str, str]]
) -> Callable[[WebDriver], Union[WebElement, bool]]:
) -> Callable[[WebDriverOrWebElement], Union[WebElement, bool]]:
"""An Expectation for checking that an element is either invisible or not
present on the DOM.
Expand All @@ -340,7 +348,7 @@ def invisibility_of_element(

def element_to_be_clickable(
mark: Union[WebElement, Tuple[str, str]]
) -> Callable[[WebDriver], Union[Literal[False], WebElement]]:
) -> Callable[[WebDriverOrWebElement], Union[Literal[False], WebElement]]:
"""An Expectation for checking an element is visible and enabled such that
you can click it.
Expand All @@ -349,7 +357,7 @@ def element_to_be_clickable(

# renamed argument to 'mark', to indicate that both locator
# and WebElement args are valid
def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
target = mark
if not isinstance(target, WebElement): # if given locator instead of WebElement
target = driver.find_element(*target) # grab element at locator
Expand All @@ -361,7 +369,7 @@ def _predicate(driver: WebDriver):
return _predicate


def staleness_of(element: WebElement) -> Callable[[WebDriver], bool]:
def staleness_of(element: WebElement) -> Callable[[Any], bool]:
"""Wait until an element is no longer attached to the DOM.
element is the element to wait for. returns False if the element is
Expand All @@ -379,7 +387,7 @@ def _predicate(_):
return _predicate


def element_to_be_selected(element: WebElement) -> Callable[[WebDriver], bool]:
def element_to_be_selected(element: WebElement) -> Callable[[Any], bool]:
"""An expectation for checking the selection is selected.
element is WebElement object
Expand All @@ -391,19 +399,19 @@ def _predicate(_):
return _predicate


def element_located_to_be_selected(locator: Tuple[str, str]) -> Callable[[WebDriver], bool]:
def element_located_to_be_selected(locator: Tuple[str, str]) -> Callable[[WebDriverOrWebElement], bool]:
"""An expectation for the element to be located is selected.
locator is a tuple of (by, path)
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
return driver.find_element(*locator).is_selected()

return _predicate


def element_selection_state_to_be(element: WebElement, is_selected: bool) -> Callable[[WebDriver], bool]:
def element_selection_state_to_be(element: WebElement, is_selected: bool) -> Callable[[Any], bool]:
"""An expectation for checking if the given element is selected.
element is WebElement object is_selected is a Boolean.
Expand All @@ -415,14 +423,16 @@ def _predicate(_):
return _predicate


def element_located_selection_state_to_be(locator: Tuple[str, str], is_selected: bool) -> Callable[[WebDriver], bool]:
def element_located_selection_state_to_be(
locator: Tuple[str, str], is_selected: bool
) -> Callable[[WebDriverOrWebElement], bool]:
"""An expectation to locate an element and check if the selection state
specified is in that state.
locator is a tuple of (by, path) is_selected is a boolean
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
try:
element = driver.find_element(*locator)
return element.is_selected() == is_selected
Expand Down Expand Up @@ -464,14 +474,14 @@ def _predicate(driver: WebDriver):
return _predicate


def element_attribute_to_include(locator: Tuple[str, str], attribute_: str) -> Callable[[WebDriver], bool]:
def element_attribute_to_include(locator: Tuple[str, str], attribute_: str) -> Callable[[WebDriverOrWebElement], bool]:
"""An expectation for checking if the given attribute is included in the
specified element.
locator, attribute
"""

def _predicate(driver: WebDriver):
def _predicate(driver: WebDriverOrWebElement):
try:
element_attribute = driver.find_element(*locator).get_attribute(attribute_)
return element_attribute is not None
Expand All @@ -481,14 +491,14 @@ def _predicate(driver: WebDriver):
return _predicate


def any_of(*expected_conditions: Callable[[WebDriver], T]) -> Callable[[WebDriver], Union[Literal[False], T]]:
def any_of(*expected_conditions: Callable[[D], T]) -> Callable[[D], Union[Literal[False], T]]:
"""An expectation that any of multiple expected conditions is true.
Equivalent to a logical 'OR'. Returns results of the first matching
condition, or False if none do.
"""

def any_of_condition(driver: WebDriver):
def any_of_condition(driver: D):
for expected_condition in expected_conditions:
try:
result = expected_condition(driver)
Expand All @@ -502,16 +512,16 @@ def any_of_condition(driver: WebDriver):


def all_of(
*expected_conditions: Callable[[WebDriver], Union[T, Literal[False]]]
) -> Callable[[WebDriver], Union[List[T], Literal[False]]]:
*expected_conditions: Callable[[D], Union[T, Literal[False]]]
) -> Callable[[D], Union[List[T], Literal[False]]]:
"""An expectation that all of multiple expected conditions is true.
Equivalent to a logical 'AND'.
Returns: When any ExpectedCondition is not met: False.
When all ExpectedConditions are met: A List with each ExpectedCondition's return value.
"""

def all_of_condition(driver: WebDriver):
def all_of_condition(driver: D):
results: List[T] = []
for expected_condition in expected_conditions:
try:
Expand All @@ -526,13 +536,13 @@ def all_of_condition(driver: WebDriver):
return all_of_condition


def none_of(*expected_conditions: Callable[[WebDriver], Any]) -> Callable[[WebDriver], bool]:
def none_of(*expected_conditions: Callable[[D], Any]) -> Callable[[D], bool]:
"""An expectation that none of 1 or multiple expected conditions is true.
Equivalent to a logical 'NOT-OR'. Returns a Boolean
"""

def none_of_condition(driver: WebDriver):
def none_of_condition(driver: D):
for expected_condition in expected_conditions:
try:
result = expected_condition(driver)
Expand Down
14 changes: 9 additions & 5 deletions py/selenium/webdriver/support/wait.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,24 +18,28 @@
import time
import typing
from typing import Callable
from typing import Generic
from typing import Literal
from typing import TypeVar
from typing import Union

from selenium.common.exceptions import NoSuchElementException
from selenium.common.exceptions import TimeoutException
from selenium.types import WaitExcTypes
from selenium.webdriver.remote.webdriver import WebDriver
from selenium.webdriver.remote.webelement import WebElement

POLL_FREQUENCY: float = 0.5 # How long to sleep in between calls to the method
IGNORED_EXCEPTIONS: typing.Tuple[typing.Type[Exception]] = (NoSuchElementException,) # default to be ignored.

T = typing.TypeVar("T")
D = TypeVar("D", bound=Union[WebDriver, WebElement])
T = TypeVar("T")


class WebDriverWait:
class WebDriverWait(Generic[D]):
def __init__(
self,
driver: WebDriver,
driver: D,
timeout: float,
poll_frequency: float = POLL_FREQUENCY,
ignored_exceptions: typing.Optional[WaitExcTypes] = None,
Expand Down Expand Up @@ -74,7 +78,7 @@ def __init__(
def __repr__(self):
return f'<{type(self).__module__}.{type(self).__name__} (session="{self._driver.session_id}")>'

def until(self, method: Callable[[WebDriver], Union[Literal[False], T]], message: str = "") -> T:
def until(self, method: Callable[[D], Union[Literal[False], T]], message: str = "") -> T:
"""Calls the method provided with the driver as an argument until the \
return value does not evaluate to ``False``.
Expand All @@ -100,7 +104,7 @@ def until(self, method: Callable[[WebDriver], Union[Literal[False], T]], message
break
raise TimeoutException(message, screen, stacktrace)

def until_not(self, method: Callable[[WebDriver], T], message: str = "") -> Union[T, Literal[True]]:
def until_not(self, method: Callable[[D], T], message: str = "") -> Union[T, Literal[True]]:
"""Calls the method provided with the driver as an argument until the \
return value evaluates to ``False``.
Expand Down

0 comments on commit 10adfe8

Please sign in to comment.