1010import shutil
1111import time
1212from typing import List
13+ from contextlib import contextmanager
1314
1415from botcity .base import BaseBot , State
1516from botcity .base .utils import only_if_element
2122from selenium .webdriver .common .by import By
2223from selenium .webdriver .common .keys import Keys
2324from selenium .webdriver .remote .webelement import WebElement
24- from selenium .webdriver .support .ui import WebDriverWait
25+ from selenium .webdriver .support .wait import WebDriverWait , TimeoutException , NoSuchElementException
26+ from selenium .webdriver .support import expected_conditions as EC
2527
2628from . import config , cv2find
2729from .browsers import BROWSER_CONFIGS , Browser
@@ -233,11 +235,16 @@ def check_driver():
233235 def stop_browser (self ):
234236 """
235237 Stops the Chrome browser and clean up the User Data Directory.
238+
239+ Warning:
240+ After invoking this method, you will need to reassign your custom options and capabilities.
236241 """
237242 if not self ._driver :
238243 return
239244 self ._driver .close ()
240245 self ._driver .quit ()
246+ self .options = None
247+ self .capabilities = None
241248 self ._driver = None
242249
243250 def set_screen_resolution (self , width = None , height = None ):
@@ -854,6 +861,26 @@ def browse(self, url):
854861 """
855862 self .navigate_to (url )
856863
864+ @contextmanager
865+ def wait_for_new_page (self , waiting_time = 10000 , activate = True ):
866+ """Context manager to wait for a new page to load and activate it.
867+
868+ Args:
869+ waiting_time (int, optional): The maximum waiting time. Defaults to 10000.
870+ activate (bool, optional): Whether or not to activate the new page. Defaults to True.
871+
872+ """
873+ tabs = self .get_tabs ()
874+ yield
875+ start_time = time .time ()
876+ while tabs == self .get_tabs ():
877+ elapsed_time = (time .time () - start_time ) * 1000
878+ if elapsed_time > waiting_time :
879+ return None
880+ time .sleep (0.1 )
881+ if activate :
882+ self .activate_tab (self .get_tabs ()[- 1 ])
883+
857884 def execute_javascript (self , code ):
858885 """
859886 Execute the given javascript code.
@@ -1032,15 +1059,19 @@ def wait_for_downloads(self, timeout: int = 120000):
10321059
10331060 wait_method = BROWSER_CONFIGS .get (self .browser ).get ("wait_for_downloads" )
10341061 # waits for all the files to be completed
1035- WebDriverWait (self ._driver , timeout / 1000 , 1 ).until (wait_method )
1062+ WebDriverWait (self ._driver , timeout / 1000.0 , 1 ).until (wait_method )
10361063
1037- def find_elements (self , selector : str , by : By = By .CSS_SELECTOR ) -> List [WebElement ]:
1064+ def find_elements (self , selector : str , by : By = By .CSS_SELECTOR ,
1065+ waiting_time = 10000 , ensure_visible : bool = True ) -> List [WebElement ]:
10381066 """Find elements using the specified selector with selector type specified by `by`.
10391067
10401068 Args:
10411069 selector (str): The selector string to be used.
10421070 by (str, optional): Selector type. Defaults to By.CSS_SELECTOR.
10431071 [See more](https://selenium-python.readthedocs.io/api.html#selenium.webdriver.common.by.By)
1072+ waiting_time (int, optional): Maximum wait time (ms) to search for a hit.
1073+ Defaults to 10000ms (10s).
1074+ ensure_visible (bool, optional): Whether to wait for the element to be visible. Defaults to True.
10441075
10451076 Returns:
10461077 List[WebElement]: List of elements found.
@@ -1054,16 +1085,36 @@ def find_elements(self, selector: str, by: By = By.CSS_SELECTOR) -> List[WebElem
10541085 ...
10551086 ```
10561087 """
1057- return self ._driver .find_elements (by , selector )
1088+ if ensure_visible :
1089+ condition = EC .visibility_of_all_elements_located
1090+ else :
1091+ condition = EC .presence_of_all_elements_located
10581092
1059- def find_element (self , selector : str , by : str = By .CSS_SELECTOR ) -> WebElement :
1093+ try :
1094+ elements = WebDriverWait (
1095+ self ._driver , timeout = waiting_time / 1000.0
1096+ ).until (
1097+ condition ((by , selector ))
1098+ )
1099+ return elements
1100+ except (TimeoutException , NoSuchElementException ) as ex :
1101+ print ("Exception on find_elements" , ex )
1102+ return None
1103+
1104+ def find_element (self , selector : str , by : str = By .CSS_SELECTOR , waiting_time = 10000 ,
1105+ ensure_visible : bool = False , ensure_clickable : bool = False ) -> WebElement :
10601106 """Find an element using the specified selector with selector type specified by `by`.
10611107 If more than one element is found, the first instance is returned.
10621108
10631109 Args:
10641110 selector (str): The selector string to be used.
10651111 by (str, optional): Selector type. Defaults to By.CSS_SELECTOR.
10661112 [See more](https://selenium-python.readthedocs.io/api.html#selenium.webdriver.common.by.By)
1113+ waiting_time (int, optional): Maximum wait time (ms) to search for a hit.
1114+ Defaults to 10000ms (10s).
1115+ ensure_visible (bool, optional): Whether to wait for the element to be visible. Defaults to False.
1116+ ensure_clickable (bool, optional): Whether to wait for the element to be clickable. Defaults to False.
1117+ If True, `ensure_clickable` takes precedence over `ensure_visible`.
10671118
10681119 Returns:
10691120 WebElement: The element found.
@@ -1079,9 +1130,47 @@ def find_element(self, selector: str, by: str = By.CSS_SELECTOR) -> WebElement:
10791130 ...
10801131 ```
10811132 """
1082- out = self .find_elements (selector = selector , by = by )
1083- if out :
1084- return out [0 ]
1133+ condition = EC .visibility_of_element_located if ensure_visible else EC .presence_of_element_located
1134+ condition = EC .element_to_be_clickable if ensure_clickable else condition
1135+
1136+ try :
1137+ element = WebDriverWait (
1138+ self ._driver , timeout = waiting_time / 1000.0
1139+ ).until (
1140+ condition ((by , selector ))
1141+ )
1142+ return element
1143+ except (TimeoutException , NoSuchElementException ):
1144+ return None
1145+
1146+ def wait_for_stale_element (self , element : WebElement , timeout : int = 10000 ):
1147+ """
1148+ Wait until the WebElement element becomes stale (outdated).
1149+
1150+ Args:
1151+ element (WebElement): The element to monitor for staleness.
1152+ timeout (int, optional): Timeout in millis. Defaults to 120000.
1153+ """
1154+ try :
1155+ WebDriverWait (self ._driver , timeout = timeout / 1000.0 ).until (EC .staleness_of (element ))
1156+ except (TimeoutException , NoSuchElementException ):
1157+ pass
1158+
1159+ def wait_for_element_visibility (self , element : WebElement , visible : bool = True , waiting_time = 10000 ):
1160+ """Wait for the element to be visible or hidden.
1161+
1162+ Args:
1163+ element (WebElement): The element to wait for.
1164+ visible (bool, optional): Whether to wait for the element to be visible. Defaults to True.
1165+ waiting_time (int, optional): Maximum wait time (ms) to search for a hit.
1166+ Defaults to 10000ms (10s).
1167+ """
1168+ if visible :
1169+ wait_method = EC .visibility_of
1170+ else :
1171+ wait_method = EC .invisibility_of_element
1172+
1173+ WebDriverWait (self ._driver , timeout = waiting_time / 1000.0 ).until (wait_method (element ))
10851174
10861175 def set_file_input_element (self , element : WebElement , filepath : str ):
10871176 """Configure the filepath for upload in a file element.
0 commit comments