1+ from typing import Union
2+
3+ from selenium import webdriver
14from selenium .common .exceptions import NoSuchElementException , TimeoutException
25from selenium .webdriver .common .by import By
36from selenium .webdriver .support import expected_conditions as EC
47from selenium .webdriver .support .ui import WebDriverWait
58
6- from s_tool .utils . driver_exceptions import SToolException
9+ from s_tool .exceptions import SToolException
710
811
9- def get_session (driver ) :
12+ def get_session (driver : webdriver ) -> str :
1013 """Return Selenium Driver session id"""
1114 return driver .session_id
1215
1316
14- def visit (driver , url ) :
17+ def visit (driver : webdriver , url : str ) -> None :
1518 """visit given url"""
1619 driver .get (url )
1720
1821
19- def page_source (driver ) :
22+ def page_source (driver : webdriver ) -> str :
2023 """Return html page source"""
2124 return driver .page_source
2225
2326
24- def current_url (driver ) :
27+ def current_url (driver : webdriver ) -> str :
2528 """Return current url"""
2629 return driver .current_url
2730
2831
29- def get_locator (locator_type , locator_text ) :
32+ def get_locator (locator_type : str , locator_text : str ) -> tuple :
3033 """Return element locator
3134
3235 Args:
@@ -40,7 +43,9 @@ def get_locator(locator_type, locator_text):
4043 return getattr (By , locator ), locator_text
4144
4245
43- def get_element (driver , locator_type , locator_text , many = None ):
46+ def get_element (
47+ driver : webdriver , locator_type : str , locator_text : str , many : bool = None
48+ ):
4449 """Get element using locator type and locator text
4550
4651 Args:
@@ -72,7 +77,9 @@ def get_element(driver, locator_type, locator_text, many=None):
7277 raise SToolException ("INVALID_SELECTOR" )
7378
7479
75- def click (driver , locator_type , locator_text , click_time = 10 ):
80+ def click (
81+ driver : webdriver , locator_type : str , locator_text : str , click_time : int = 10
82+ ) -> Union [bool , None ]:
7683 """Return True if element clicked otherwise return None
7784
7885 Args:
@@ -99,7 +106,7 @@ def click(driver, locator_type, locator_text, click_time=10):
99106 raise SToolException (ex )
100107
101108
102- def get_cookies (driver ) :
109+ def get_cookies (driver : webdriver ) -> dict :
103110 """Accept driver object and return cookies in dictionary
104111
105112 Args:
@@ -114,7 +121,7 @@ def get_cookies(driver):
114121 return cookies_dict or {}
115122
116123
117- def take_screenshot (driver , element = None ):
124+ def take_screenshot (driver : webdriver , element : tuple = None ) -> Union [ bytes , None ] :
118125 """take screenshot of given element if element is
119126 not given take a full page screeenshot and return
120127 data in bytes
@@ -143,7 +150,7 @@ def take_screenshot(driver, element=None):
143150 return driver .get_screenshot_as_png ()
144151
145152
146- def display_element (driver , element , hide = None ):
153+ def display_element (driver : webdriver , element , hide = None ) -> None :
147154 """hide or show single element
148155
149156 Args:
@@ -159,7 +166,7 @@ def display_element(driver, element, hide=None):
159166 driver .execute_script (f"arguments[0].style.display = '{ hide_or_show } ';" , element )
160167
161168
162- def hide_show_elements (driver , elements , hide = None ):
169+ def hide_show_elements (driver : webdriver , elements : list , hide : bool = None ) -> None :
163170 """hide or show multiple elements
164171
165172 Args:
@@ -174,7 +181,7 @@ def hide_show_elements(driver, elements, hide=None):
174181 """
175182 for element_locator in elements :
176183 locator_type , locator_value = element_locator
177- element_list = get_element (driver , locator_type , locator_value , 1 )
184+ element_list = get_element (driver , locator_type , locator_value , many = True )
178185 if element_list :
179186 for element in element_list :
180187 display_element (driver , element , hide )
0 commit comments