@@ -49,15 +49,29 @@ def wait_for_idle_status(driver, current_cell, start_time, timeout):
4949
5050 print (current_cell .text )
5151
52+ def rewind_to_first_cell (cell ):
53+ while True :
54+ prev = cell .find_elements (
55+ By .XPATH ,
56+ "preceding-sibling::div[contains(@class,'jp-Notebook-cell')][1]"
57+ )
58+ if not prev :
59+ return cell
60+ cell = prev [0 ]
61+
62+
5263
5364def run_notebook (driver , notebook_area , args ):
5465 """This functions runs all the cells of the notebook"""
5566 print ("Running Cells" )
5667 start_time = time .time ()
68+ current_cell = driver .find_element (
69+ By .CSS_SELECTOR , ".jp-Notebook-cell.jp-mod-selected"
70+ )
71+
72+ current_cell = rewind_to_first_cell (current_cell )
73+
5774 while True :
58- current_cell = driver .find_element (
59- By .CSS_SELECTOR , ".jp-Notebook-cell.jp-mod-selected"
60- )
6175 editor_divs = current_cell .find_elements (
6276 By .CSS_SELECTOR , ".jp-InputArea-editor div"
6377 )
@@ -74,6 +88,9 @@ def run_notebook(driver, notebook_area, args):
7488
7589 if cell_is_waiting_for_input (driver ):
7690 print ("Cell requesting input" )
91+ current_cell = driver .find_element (
92+ By .CSS_SELECTOR , ".jp-Notebook-cell.jp-mod-selected"
93+ )
7794 input_box = WebDriverWait (driver , 5 ).until (
7895 EC .visibility_of_element_located ((By .CSS_SELECTOR , ".jp-Stdin-input" ))
7996 )
@@ -112,7 +129,11 @@ def run_notebook(driver, notebook_area, args):
112129 time .sleep (0.1 )
113130 if not cell_is_waiting_for_input (driver ):
114131 wait_for_idle_status (driver , current_cell , start_time , args .timeout )
115-
132+ next_cell = current_cell .find_element (
133+ By .XPATH ,
134+ "following-sibling::div[contains(@class,'jp-Notebook-cell')][1]" ,
135+ )
136+ current_cell = next_cell
116137
117138def download_notebook (driver ):
118139 """This function is used to download the notebook currently open."""
0 commit comments