client.py

from contextlib import contextmanager
from inspect import getsource
from textwrap import dedent
from typing import Any, Dict, List, Optional, Union

from nbclient import NotebookClient
from nbclient.exceptions import CellExecutionError
from nbformat.v4 import new_code_cell

from .exceptions import (
    TestbookCellTagNotFoundError,
    TestbookExecuteResultNotFoundError,
    TestbookSerializeError,
    TestbookRuntimeError,
    TestbookError,
)
from .reference import TestbookObjectReference
from .testbooknode import TestbookNode
from .translators import PythonTranslator
from .utils import random_varname, all_subclasses


class TestbookNotebookClient(NotebookClient):
    __test__ = False

    def __init__(self, nb, km=None, **kw):
        # Fix the ipykernel 5.5 issue where execute requests after errors are aborted
        ea = kw.get('extra_arguments', [])
        if not any(arg.startswith('--Kernel.stop_on_error_timeout') for arg in self.extra_arguments):
            ea.append('--Kernel.stop_on_error_timeout=0')
        kw['extra_arguments'] = ea
        super().__init__(nb, km=km, **kw)

    def ref(self, name: str) -> Union[TestbookObjectReference, Any]:
        """
        Return a reference to an object in the kernel
        """

        # Check if exists
        self.inject(name, pop=True)
        try:
            self.inject(f"import json; json.dumps({name})", pop=True)
            return self.value(name)
        except Exception:
            return TestbookObjectReference(self, name)

    def get(self, item):
        return self.ref(item)

    def __getitem__(self, item):
        return self.ref(item)

    @staticmethod
    def _construct_call_code(
        func_name: str, args: Optional[List] = None, kwargs: Optional[Dict] = None
    ) -> str:
        return """
            {func_name}(*{args_list}, **{kwargs_dict})
            """.format(
            func_name=func_name,
            args_list=PythonTranslator.translate(args) if args else [],
            kwargs_dict=PythonTranslator.translate(kwargs) if kwargs else {},
        )

    @property
    def cells(self):
        return self.nb.cells

    @staticmethod
    def _execute_result(cell) -> List:
        """
        Return data from execute_result outputs
        """

        return [
            output["data"]
            for output in cell["outputs"]
            if output["output_type"] == 'execute_result'
        ]

    @staticmethod
    def _output_text(cell) -> str:
        if "outputs" not in cell:
            raise ValueError("cell must be a code cell")

        text = ''
        for output in cell["outputs"]:
            if 'text' in output:
                text += output['text']
            elif "data" in output and "text/plain" in output["data"]:
                text += output["data"]["text/plain"]

        return text.strip()

    def _cell_index(self, tag: Union[int, str]) -> int:
        """
        Get cell index from the cell tag
        """

        if isinstance(tag, int):
            return tag
        elif not isinstance(tag, str):
            raise TypeError('expected tag as str')

        for idx, cell in enumerate(self.cells):
            metadata = cell['metadata']
            if "tags" in metadata and tag in metadata['tags']:
                return idx

        raise TestbookCellTagNotFoundError("Cell tag '{}' not found".format(tag))

    def execute_cell(self, cell, **kwargs) -> Union[Dict, List[Dict]]:
        """
        Executes a cell or list of cells
        """
        if isinstance(cell, slice):
            start, stop = self._cell_index(cell.start), self._cell_index(cell.stop)
            if cell.step is not None:
                raise TestbookError('testbook does not support step argument')

            cell = range(start, stop + 1)
        elif isinstance(cell, str) or isinstance(cell, int):
            cell = [cell]

        cell_indexes = cell

        if all(isinstance(x, str) for x in cell):
            cell_indexes = [self._cell_index(tag) for tag in cell]

        executed_cells = []
        for idx in cell_indexes:
            try:
                cell = super().execute_cell(self.nb['cells'][idx], idx, **kwargs)
            except CellExecutionError as ce:
                raise TestbookRuntimeError(ce.evalue, ce, self._get_error_class(ce.ename))

            executed_cells.append(cell)

        return executed_cells[0] if len(executed_cells) == 1 else executed_cells

    def execute(self) -> None:
        """
        Executes all cells
        """

        for index, cell in enumerate(self.nb.cells):
            super().execute_cell(cell, index)

    def cell_output_text(self, cell) -> str:
        """
        Return cell text output
        """
        cell_index = self._cell_index(cell)
        return self._output_text(self.nb['cells'][cell_index])

    def cell_execute_result(self, cell: Union[int, str]) -> List[Dict[str, Any]]:
        """Return the execute results of cell at a given index or with a given tag.

        Each result is expressed with a dictionary for which the key is the mimetype
        of the data. A same result can have different representation corresponding to
        different mimetype.

        Parameters
        ----------
        cell : int or str
            The index or tag to look for

        Returns
        -------
        List[Dict[str, Any]]
            The execute results

        Raises
        ------
        IndexError
            If index is invalid
        TestbookCellTagNotFoundError
            If tag is not found
        """
        cell_index = self._cell_index(cell)
        return self._execute_result(self.nb['cells'][cell_index])

    def inject(
        self,
        code: str,
        args: List = None,
        kwargs: Dict = None,
        run: bool = True,
        before: Optional[Union[str, int]] = None,
        after: Optional[Union[str, int]] = None,
        pop: bool = False,
    ) -> TestbookNode:
        """Injects and executes given code block

        Parameters
        ----------
        code : str
            Code or function to be injected
        args : iterable, optional
            tuple of arguments to be passed to the function
        kwargs : dict, optional
            dict of keyword arguments to be passed to the function
        run : bool, optional
            Control immediate execution after injection (default is True)
        before, after : int, str, optional
            Inject code before or after cell
        pop : bool
            Pop cell after execution (default is False)

        Returns
        -------
        TestbookNode
            Injected cell
        """

        if isinstance(code, str):
            lines = dedent(code)
        elif callable(code):
            lines = getsource(code) + (
                dedent(self._construct_call_code(code.__name__, args, kwargs)) if run else ''
            )
        else:
            raise TypeError('can only inject function or code block as str')

        inject_idx = len(self.cells)

        if after is not None and before is not None:
            raise ValueError("pass either before or after as kwargs")
        elif before is not None:
            inject_idx = self._cell_index(before)
        elif after is not None:
            inject_idx = self._cell_index(after) + 1

        code_cell = new_code_cell(lines)
        self.cells.insert(inject_idx, code_cell)

        cell = TestbookNode(self.execute_cell(inject_idx)) if run else TestbookNode(code_cell)

        if self._contains_error(cell):
            eclass = self._get_error_class(cell.get('outputs')[0]['ename'])
            evalue = cell.get('outputs')[0]['evalue']
            raise TestbookRuntimeError(evalue, evalue, eclass)

        if run and pop:
            self.cells.pop(inject_idx)

        return cell

    def value(self, code: str) -> Any:
        """
        Execute given code in the kernel and return JSON serializeable result.

        If the result is not JSON serializeable, it raises `TestbookAttributeError`.
        This error object will also contain an attribute called `save_varname` which
        can be used to create a reference object with :meth:`ref`.

        Parameters
        ----------
        code: str
            This can be any executable code that returns a value.
            It can be used the return the value of an object, or the output
            of a function call.

        Returns
        -------
        The output of the executed code

        Raises
        ------
            TestbookSerializeError

        """
        result = self.inject(code, pop=True)

        if not self._execute_result(result):
            raise TestbookExecuteResultNotFoundError(
                'code provided does not produce execute_result'
            )

        save_varname = random_varname()

        inject_code = f"""
            import json
            from IPython import get_ipython
            from IPython.display import JSON

            {save_varname} = get_ipython().last_execution_result.result

            json.dumps({save_varname})
            JSON({{"value" : {save_varname}}})
        """

        try:
            outputs = self.inject(inject_code, pop=True).outputs

            if outputs[0].output_type == "error":
                # will receive error when `allow_errors` is set to True
                raise TestbookRuntimeError(
                    outputs[0].evalue, outputs[0].traceback, outputs[0].ename
                )

            return outputs[0].data['application/json']['value']

        except TestbookRuntimeError:
            e = TestbookSerializeError('could not JSON serialize output')
            e.save_varname = save_varname
            raise e

    @contextmanager
    def patch(self, target, **kwargs):
        """Used as contextmanager to patch objects in the kernel"""
        mock_object = f'_mock_{random_varname()}'
        patcher = f'_patcher_{random_varname()}'

        self.inject(
            f"""
            from unittest.mock import patch
            {patcher} = patch(
                {PythonTranslator.translate(target)},
                **{PythonTranslator.translate(kwargs)}
            )
            {mock_object} = {patcher}.start()
        """
        )

        yield TestbookObjectReference(self, mock_object)

        self.inject(f"{patcher}.stop()")

    @contextmanager
    def patch_dict(self, in_dict, values=(), clear=False, **kwargs):
        """Used as contextmanager to patch dictionaries in the kernel"""
        mock_object = f'_mock_{random_varname()}'
        patcher = f'_patcher_{random_varname()}'

        self.inject(
            f"""
            from unittest.mock import patch
            {patcher} = patch.dict(
                {PythonTranslator.translate(in_dict)},
                {PythonTranslator.translate(values)},
                {PythonTranslator.translate(clear)},
                **{PythonTranslator.translate(kwargs)}
            )
            {mock_object} = {patcher}.start()
        """
        )

        yield TestbookObjectReference(self, mock_object)

        self.inject(f"{patcher}.stop()")

    @staticmethod
    def _get_error_class(ename):
        eclass = None
        for klass in all_subclasses(Exception):
            if klass.__name__ == ename:
                eclass = klass
                break
        return eclass

    @staticmethod
    def _contains_error(result):
        return result.get('outputs') and result.get('outputs')[0].output_type == "error"