[py] support for custom error messages through functions for WebDriverWait.until/until_not

Author: Delta456

py/selenium/webdriver/support/wait.py

@@ -16,6 +16,7 @@
 # under the License.
 
 import time
+from typing import Any
 from typing import Callable
 from typing import Generic
 from typing import Literal
@@ -92,7 +93,9 @@ def __init__(
     def __repr__(self) -> str:
         return f'<{type(self).__module__}.{type(self).__name__} (session="{self._driver.session_id}")>'
 
-    def until(self, method: Callable[[D], Union[Literal[False], T]], message: str = "") -> T:
+    def until(
+        self, method: Callable[[D], Union[Literal[False], T]], message: Union[str, Callable[[Any], str]] = ""
+    ) -> T:
         """Wait until the method returns a value that is not False.
 
         Calls the method provided with the driver as an argument until the
@@ -103,7 +106,7 @@ def until(self, method: Callable[[D], Union[Literal[False], T]], message: str =
         method: callable(WebDriver)
             - A callable object that takes a WebDriver instance as an argument.
 
-        message: str
+        message: Union[str, Callable[[Any], str]]
             - Optional message for :exc:`TimeoutException`
 
         Return:
@@ -143,9 +146,13 @@ def until(self, method: Callable[[D], Union[Literal[False], T]], message: str =
             if time.monotonic() > end_time:
                 break
             time.sleep(self._poll)
-        raise TimeoutException(message, screen, stacktrace)
 
-    def until_not(self, method: Callable[[D], T], message: str = "") -> Union[T, Literal[True]]:
+        final_msg = message() if callable(message) else message
+        raise TimeoutException(final_msg, screen, stacktrace)
+
+    def until_not(
+        self, method: Callable[[D], T], message: Union[str, Callable[[Any], str]] = ""
+    ) -> Union[T, Literal[True]]:
         """Wait until the method returns a value that is not False.
 
         Calls the method provided with the driver as an argument until the
@@ -156,7 +163,7 @@ def until_not(self, method: Callable[[D], T], message: str = "") -> Union[T, Lit
         method: callable(WebDriver)
             - A callable object that takes a WebDriver instance as an argument.
 
-        message: str
+        message: Union[str, Callable[[Any], str]]
             - Optional message for :exc:`TimeoutException`
 
         Return:
@@ -192,4 +199,5 @@ def until_not(self, method: Callable[[D], T], message: str = "") -> Union[T, Lit
             if time.monotonic() > end_time:
                 break
             time.sleep(self._poll)
-        raise TimeoutException(message)
+        final_msg = message() if callable(message) else message
+        raise TimeoutException(final_msg)

py/test/selenium/webdriver/common/webdriverwait_tests.py

@@ -58,6 +58,21 @@ def test_should_still_fail_to_find_an_element_with_explicit_wait(driver, pages):
         WebDriverWait(driver, 0.01).until(EC.presence_of_element_located((By.ID, "box0")))
 
 
+def test_should_still_fail_to_find_an_element_with_explicit_wait_with_custom_timeout_messages(driver, pages):
+    pages.load("dynamic.html")
+    with pytest.raises(TimeoutException) as exception:
+        WebDriverWait(driver, 0.01).until(
+            EC.presence_of_element_located((By.ID, "box0")), message=lambda: "custom timeout message"
+        )
+    assert "custom timeout message" in str(exception.value)
+
+    with pytest.raises(TimeoutException) as exception:
+        WebDriverWait(driver, 0.01).until(
+            EC.presence_of_element_located((By.ID, "box0")), message="custom timeout message"
+        )
+    assert "custom timeout message" in str(exception.value)
+
+
 def test_should_explicitly_wait_until_at_least_one_element_is_found_when_searching_for_many(driver, pages):
     pages.load("dynamic.html")
     add = driver.find_element(By.ID, "adder")