import os import time from selenium.webdriver import ActionChains from selenium.webdriver.common.by import By from selenium.webdriver.common.keys import Keys from selenium.webdriver.support.ui import WebDriverWait from selenium.webdriver.support import expected_conditions as EC from selenium.webdriver.remote.webelement import WebElement from contextlib import contextmanager pjoin = os.path.join def wait_for_selector(browser, selector, timeout=10, visible=False, single=False): wait = WebDriverWait(browser, timeout) if single: if visible: conditional = EC.visibility_of_element_located else: conditional = EC.presence_of_element_located else: if visible: conditional = EC.visibility_of_all_elements_located else: conditional = EC.presence_of_all_elements_located return wait.until(conditional((By.CSS_SELECTOR, selector))) class CellTypeError(ValueError): def __init__(self, message=""): self.message = message class Notebook: def __init__(self, browser): self.browser = browser self.disable_autosave_and_onbeforeunload() def __len__(self): return len(self.cells) def __getitem__(self, key): return self.cells[key] def __setitem__(self, key, item): if isinstance(key, int): self.edit_cell(index=key, content=item, render=False) # TODO: re-add slicing support, handle general python slicing behaviour # includes: overwriting the entire self.cells object if you do # self[:] = [] # elif isinstance(key, slice): # indices = (self.index(cell) for cell in self[key]) # for k, v in zip(indices, item): # self.edit_cell(index=k, content=v, render=False) def __iter__(self): return (cell for cell in self.cells) @property def body(self): return self.browser.find_element_by_tag_name("body") @property def cells(self): """Gets all cells once they are visible. """ return self.browser.find_elements_by_class_name("cell") @property def current_index(self): return self.index(self.current_cell) def index(self, cell): return self.cells.index(cell) def disable_autosave_and_onbeforeunload(self): """Disable request to save before closing window and autosave. This is most easily done by using js directly. """ self.browser.execute_script("window.onbeforeunload = null;") self.browser.execute_script("Jupyter.notebook.set_autosave_interval(0)") def to_command_mode(self): """Changes us into command mode on currently focused cell """ self.body.send_keys(Keys.ESCAPE) self.browser.execute_script("return Jupyter.notebook.handle_command_mode(" "Jupyter.notebook.get_cell(" "Jupyter.notebook.get_edit_index()))") def focus_cell(self, index=0): cell = self.cells[index] cell.click() self.to_command_mode() self.current_cell = cell def find_and_replace(self, index=0, find_txt='', replace_txt=''): self.focus_cell(index) self.to_command_mode() self.body.send_keys('f') wait_for_selector(self.browser, "#find-and-replace", single=True) self.browser.find_element_by_id("findreplace_allcells_btn").click() self.browser.find_element_by_id("findreplace_find_inp").send_keys(find_txt) self.browser.find_element_by_id("findreplace_replace_inp").send_keys(replace_txt) self.browser.find_element_by_id("findreplace_replaceall_btn").click() def convert_cell_type(self, index=0, cell_type="code"): # TODO add check to see if it is already present self.focus_cell(index) cell = self.cells[index] if cell_type == "markdown": self.current_cell.send_keys("m") elif cell_type == "raw": self.current_cell.send_keys("r") elif cell_type == "code": self.current_cell.send_keys("y") else: raise CellTypeError(("{} is not a valid cell type," "use 'code', 'markdown', or 'raw'").format(cell_type)) self.wait_for_stale_cell(cell) self.focus_cell(index) return self.current_cell def wait_for_stale_cell(self, cell): """ This is needed to switch a cell's mode and refocus it, or to render it. Warning: there is currently no way to do this when changing between markdown and raw cells. """ wait = WebDriverWait(self.browser, 10) element = wait.until(EC.staleness_of(cell)) def get_cells_contents(self): JS = 'return Jupyter.notebook.get_cells().map(function(c) {return c.get_text();})' return self.browser.execute_script(JS) def get_cell_contents(self, index=0, selector='div .CodeMirror-code'): return self.cells[index].find_element_by_css_selector(selector).text def set_cell_metadata(self, index, key, value): JS = 'Jupyter.notebook.get_cell({}).metadata.{} = {}'.format(index, key, value) return self.browser.execute_script(JS) def get_cell_type(self, index=0): JS = 'return Jupyter.notebook.get_cell({}).cell_type'.format(index) return self.browser.execute_script(JS) def set_cell_input_prompt(self, index, prmpt_val): JS = 'Jupyter.notebook.get_cell({}).set_input_prompt({})'.format(index, prmpt_val) self.browser.execute_script(JS) def edit_cell(self, cell=None, index=0, content="", render=False): """Set the contents of a cell to *content*, by cell object or by index """ if cell is not None: index = self.index(cell) self.focus_cell(index) # Select & delete anything already in the cell self.current_cell.send_keys(Keys.ENTER) ctrl(self.browser, 'a') self.current_cell.send_keys(Keys.DELETE) for line_no, line in enumerate(content.splitlines()): if line_no != 0: self.current_cell.send_keys(Keys.ENTER, "\n") self.current_cell.send_keys(Keys.ENTER, line) if render: self.execute_cell(self.current_index) def execute_cell(self, cell_or_index=None): if isinstance(cell_or_index, int): index = cell_or_index elif isinstance(cell_or_index, WebElement): index = self.index(cell_or_index) else: raise TypeError("execute_cell only accepts a WebElement or an int") self.focus_cell(index) self.current_cell.send_keys(Keys.CONTROL, Keys.ENTER) def add_cell(self, index=-1, cell_type="code", content=""): self.focus_cell(index) self.current_cell.send_keys("b") new_index = index + 1 if index >= 0 else index if content: self.edit_cell(index=index, content=content) if cell_type != 'code': self.convert_cell_type(index=new_index, cell_type=cell_type) def delete_cell(self, index): self.focus_cell(index) self.to_command_mode() self.current_cell.send_keys('dd') def add_markdown_cell(self, index=-1, content="", render=True): self.add_cell(index, cell_type="markdown") self.edit_cell(index=index, content=content, render=render) def append(self, *values, cell_type="code"): for i, value in enumerate(values): if isinstance(value, str): self.add_cell(cell_type=cell_type, content=value) else: raise TypeError("Don't know how to add cell from %r" % value) def extend(self, values): self.append(*values) def run_all(self): for cell in self: self.execute_cell(cell) def trigger_keydown(self, keys): trigger_keystrokes(self.body, keys) @classmethod def new_notebook(cls, browser, kernel_name='kernel-python3'): with new_window(browser, selector=".cell"): select_kernel(browser, kernel_name=kernel_name) return cls(browser) def select_kernel(browser, kernel_name='kernel-python3'): """Clicks the "new" button and selects a kernel from the options. """ wait = WebDriverWait(browser, 10) new_button = wait.until(EC.element_to_be_clickable((By.ID, "new-dropdown-button"))) new_button.click() kernel_selector = '#{} a'.format(kernel_name) kernel = wait_for_selector(browser, kernel_selector, single=True) kernel.click() @contextmanager def new_window(browser, selector=None): """Contextmanager for switching to & waiting for a window created. This context manager gives you the ability to create a new window inside the created context and it will switch you to that new window. If you know a CSS selector that can be expected to appear on the window, then this utility can wait on that selector appearing on the page before releasing the context. Usage example: from notebook.tests.selenium.utils import new_window, Notebook ⋮ # something that creates a browser object with new_window(browser, selector=".cell"): select_kernel(browser, kernel_name=kernel_name) nb = Notebook(browser) """ initial_window_handles = browser.window_handles yield new_window_handle = next(window for window in browser.window_handles if window not in initial_window_handles) browser.switch_to_window(new_window_handle) if selector is not None: wait_for_selector(browser, selector) def shift(browser, k): """Send key combination Shift+(k)""" trigger_keystrokes(browser, "shift-%s"%k) def ctrl(browser, k): """Send key combination Ctrl+(k)""" trigger_keystrokes(browser, "control-%s"%k) def trigger_keystrokes(browser, *keys): """ Send the keys in sequence to the browser. Handles following key combinations 1. with modifiers eg. 'control-alt-a', 'shift-c' 2. just modifiers eg. 'alt', 'esc' 3. non-modifiers eg. 'abc' Modifiers : http://seleniumhq.github.io/selenium/docs/api/py/webdriver/selenium.webdriver.common.keys.html """ for each_key_combination in keys: keys = each_key_combination.split('-') if len(keys) > 1: # key has modifiers eg. control, alt, shift modifiers_keys = [getattr(Keys, x.upper()) for x in keys[:-1]] ac = ActionChains(browser) for i in modifiers_keys: ac = ac.key_down(i) ac.send_keys(keys[-1]) for i in modifiers_keys[::-1]: ac = ac.key_up(i) ac.perform() else: # single key stroke. Check if modifier eg. "up" browser.send_keys(getattr(Keys, keys[0].upper(), keys[0]))