diff --git a/nbclient/client.py b/nbclient/client.py index 2d7beeb9..b1b98845 100644 --- a/nbclient/client.py +++ b/nbclient/client.py @@ -14,7 +14,18 @@ from jupyter_client.client import KernelClient from nbformat import NotebookNode from nbformat.v4 import output_from_msg -from traitlets import Any, Bool, Dict, Enum, Integer, List, Type, Unicode, default +from traitlets import ( + Any, + Bool, + Callable, + Dict, + Enum, + Integer, + List, + Type, + Unicode, + default, +) from traitlets.config.configurable import LoggingConfigurable from .exceptions import ( @@ -32,6 +43,9 @@ def timestamp() -> str: return datetime.datetime.utcnow().isoformat() + 'Z' +MD_EXPRESSIONS_PREFIX = "md-expr" + + class NotebookClient(LoggingConfigurable): """ Encompasses a Client for executing cells in a notebook @@ -285,6 +299,16 @@ def _kernel_manager_class_default(self) -> KernelManager: """, ).tag(config=True) + parse_md_expressions: t.Optional[t.Callable[[str], t.List[str]]] = Callable( + None, + allow_none=True, + help=dedent( + """ + A function to extract expression variables from a Markdown cell. + """ + ), + ) + resources: t.Dict = Dict( help=dedent( """ @@ -788,6 +812,12 @@ async def async_execute_cell( The cell which was just processed. """ assert self.kc is not None + + if self.parse_md_expressions and cell.cell_type == 'markdown': + expressions = self.parse_md_expressions(cell.source) + if expressions: + await self.async_execute_expressions(cell, cell_index, expressions) + if cell.cell_type != 'code' or not cell.source.strip(): self.log.debug("Skipping non-executing cell %s", cell_index) return cell @@ -1009,6 +1039,110 @@ def handle_comm_msg(self, outs: t.List, msg: t.Dict, cell_index: int) -> None: if comm_id in self.comm_objects: self.comm_objects[comm_id].handle_msg(msg) + async def async_execute_expressions( + self, cell, cell_index: int, expressions: t.List[str] + ) -> t.Dict[str, Any]: + parent_msg_id = await ensure_async( + self.kc.execute( + '', + silent=True, + user_expressions={ + f"{MD_EXPRESSIONS_PREFIX}-{i}": expr for i, expr in enumerate(expressions) + }, + ) + ) + task_poll_kernel_alive = asyncio.ensure_future(self._async_poll_kernel_alive()) + task_poll_expr_msg = asyncio.ensure_future(self._async_poll_expr_msg(parent_msg_id)) + exec_timeout = None + self.task_poll_for_reply = asyncio.ensure_future( + self._async_poll_for_expr_reply( + parent_msg_id, cell, exec_timeout, task_poll_expr_msg, task_poll_kernel_alive + ) + ) + try: + exec_reply = await self.task_poll_for_reply + except asyncio.CancelledError: + # can only be cancelled by task_poll_kernel_alive when the kernel is dead + task_poll_expr_msg.cancel() + raise DeadKernelError("Kernel died") + except Exception as e: + # Best effort to cancel request if it hasn't been resolved + try: + # Check if the task_poll_output is doing the raising for us + if not isinstance(e, CellControlSignal): + task_poll_expr_msg.cancel() + finally: + raise + self._check_raise_for_error(cell, exec_reply) + attachments = { + key: val["data"] if "data" in val else {"traceback": "\n".join(val["traceback"])} + for key, val in exec_reply["content"]["user_expressions"].items() + } + cell.setdefault("attachments", {}) + # remove old expressions from cell + cell["attachments"] = { + key: val + for key, val in cell["attachments"].items() + if not key.startswith(MD_EXPRESSIONS_PREFIX) + } + cell["attachments"].update(attachments) + self.nb['cells'][cell_index] = cell + return cell + + async def _async_poll_for_expr_reply( + self, + msg_id: str, + cell: NotebookNode, + timeout: t.Optional[int], + task_poll_output_msg: asyncio.Future, + task_poll_kernel_alive: asyncio.Future, + ) -> t.Dict: + + assert self.kc is not None + new_timeout: t.Optional[float] = None + if timeout is not None: + deadline = monotonic() + timeout + new_timeout = float(timeout) + while True: + try: + msg = await ensure_async(self.kc.shell_channel.get_msg(timeout=new_timeout)) + if msg['parent_header'].get('msg_id') == msg_id: + try: + await asyncio.wait_for(task_poll_output_msg, self.iopub_timeout) + except (asyncio.TimeoutError, Empty): + if self.raise_on_iopub_timeout: + task_poll_kernel_alive.cancel() + raise CellTimeoutError.error_from_timeout_and_cell( + "Timeout waiting for IOPub output", self.iopub_timeout, cell + ) + else: + self.log.warning("Timeout waiting for IOPub output") + task_poll_kernel_alive.cancel() + return msg + else: + if new_timeout is not None: + new_timeout = max(0, deadline - monotonic()) + except Empty: + # received no message, check if kernel is still alive + assert timeout is not None + task_poll_kernel_alive.cancel() + await self._async_check_alive() + await self._async_handle_timeout(timeout, cell) + + async def _async_poll_expr_msg(self, parent_msg_id: str) -> None: + + assert self.kc is not None + while True: + msg = await ensure_async(self.kc.iopub_channel.get_msg(timeout=None)) + if msg['parent_header'].get('msg_id') == parent_msg_id: + try: + msg_type = msg['msg_type'] + if msg_type == 'status': + if msg['content']['execution_state'] == 'idle': + raise CellExecutionComplete() + except CellExecutionComplete: + return + def _serialize_widget_state(self, state: t.Dict) -> t.Dict[str, t.Any]: """Serialize a widget state, following format in @jupyter-widgets/schema.""" return { diff --git a/nbclient/tests/files/Markdown_expressions.ipynb b/nbclient/tests/files/Markdown_expressions.ipynb new file mode 100644 index 00000000..45af6df0 --- /dev/null +++ b/nbclient/tests/files/Markdown_expressions.ipynb @@ -0,0 +1,31 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "a = 1\n", + "b = 'c'" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "attachments": { + "md-expr-0": {"text/plain": "1"}, + "md-expr-1": {"text/plain": "'c'"}, + "md-expr-2": {"traceback": "\u001b[0;31mNameError\u001b[0m\u001b[0;31m:\u001b[0m name 'c' is not defined\n"} + }, + "source": [ + "# Variables\n", + "\n", + "{{ a }} {{ b }} {{ c }}" + ] + } + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/nbclient/tests/test_client.py b/nbclient/tests/test_client.py index 6d109f19..ff099692 100644 --- a/nbclient/tests/test_client.py +++ b/nbclient/tests/test_client.py @@ -230,6 +230,10 @@ def assert_notebooks_equal(expected, actual): actual_execution_count = actual_cell.get('execution_count', None) assert expected_execution_count == actual_execution_count + expected_execution_count = expected_cell.get('attachments', None) + actual_execution_count = actual_cell.get('attachments', None) + assert expected_execution_count == actual_execution_count + def notebook_resources(): """ @@ -249,6 +253,23 @@ def filter_messages_on_error_output(err_output): return os.linesep.join(filtered_result) +def parse_md_expressions(source: str) -> list: + """ + Parse markdown expressions from a string. + + :param source: The source string to parse. + :return: A list of markdown expressions. + """ + from markdown_it import MarkdownIt, tree + from mdit_py_plugins.substitution import substitution_plugin + + mdit = MarkdownIt().use(substitution_plugin) + tokens = tree.SyntaxTreeNode(mdit.parse(source)) + return [ + t.content for t in tokens.walk() if t.type in ['substitution_inline', 'substitution_block'] + ] + + @pytest.mark.parametrize( ["input_name", "opts"], [ @@ -271,6 +292,10 @@ def filter_messages_on_error_output(err_output): ("UnicodePy3.ipynb", dict(kernel_name="python")), ("update-display-id.ipynb", dict(kernel_name="python")), ("Check History in Memory.ipynb", dict(kernel_name="python")), + ( + "Markdown_expressions.ipynb", + dict(kernel_name="python", parse_md_expressions=parse_md_expressions), + ), ], ) def test_run_all_notebooks(input_name, opts): diff --git a/requirements-dev.txt b/requirements-dev.txt index a9b9e078..f75987e4 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -15,3 +15,4 @@ pip>=18.1 wheel>=0.31.0 setuptools>=38.6.0 twine>=1.11.0 +markdown-it-py[plugins]~=1.0.0