feat: add ability to run commands after a cell (#94)

* add post-command
* bump isort to fix Poetry configuration is invalid
* lints
* format post command in error message
* handle masked exceptions
* s/post_command/post_cell_execute/
* s/post command/post cell execution

Author: tekumara

.pre-commit-config.yaml

@@ -19,7 +19,7 @@ repos:
       - id: black
         files: '(src|tests).*py'
   - repo: https://github.com/PyCQA/isort
-    rev: 5.6.4
+    rev: 5.12.0
     hooks:
       - id: isort
         args: ["-m", "3", "--tc"]

src/nbmake/nb_run.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 import nbformat
 from nbclient.client import (
@@ -79,20 +79,34 @@ async def apply_mocks(
                         if execute_reply["content"]["ename"] == "ModuleNotFoundError":
                             if self.find_import_errors:
                                 raise CellImportError()
-
-                if c.kc is None:
-                    raise Exception("there is no kernelclient")
-                mocks: Dict[str, Any] = (
-                    cell.get("metadata", {}).get("nbmake", {}).get("mock", {})
-                )
-                for v in mocks:
-                    if isinstance(mocks[v], str):
-                        out = await c.kc.execute_interactive(f"{v} = '{mocks[v]}'")
-                    else:
-                        out = await c.kc.execute_interactive(f"{v} = {mocks[v]}")
-
-                    if out["content"]["status"] != "ok":
-                        raise Exception(f"Failed to apply mock {v}\n\n{str(out)}")
+                else:
+                    if c.kc is None:
+                        raise Exception("there is no kernelclient")
+                    mocks: Dict[str, Any] = (
+                        cell.get("metadata", {}).get("nbmake", {}).get("mock", {})
+                    )
+                    for v in mocks:
+                        if isinstance(mocks[v], str):
+                            out = await c.kc.execute_interactive(f"{v} = '{mocks[v]}'")
+                        else:
+                            out = await c.kc.execute_interactive(f"{v} = {mocks[v]}")
+
+                        if out["content"]["status"] != "ok":
+                            raise Exception(f"Failed to apply mock {v}\n\n{str(out)}")
+
+                    post_cell_execute: List[str] = (
+                        cell.get("metadata", {})
+                        .get("nbmake", {})
+                        .get("post_cell_execute", [])
+                    )
+                    if post_cell_execute:
+                        pce = "\n".join(post_cell_execute)
+                        out = await c.kc.execute_interactive(pce)
+
+                        if out["content"]["status"] != "ok":
+                            raise Exception(
+                                f"Post cell execution failed:\n{pce}\n\n{str(out)}"
+                            )
 
             c.on_cell_executed = apply_mocks
 

tests/resources/post_cell_execute.ipynb

@@ -0,0 +1,51 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "nbmake": {
+     "post_cell_execute": [
+      "y = 3",
+      "z = x+y"
+     ]
+    }
+   },
+   "outputs": [],
+   "source": [
+    "x = 1\n",
+    "y = 2\n",
+    "z = 0\n",
+    "# this cell has a post_cell_execute that assigns y and z"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert y == 3\n",
+    "assert z == 4"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.5"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "b62bdd1a52024cb952ae76a82719d36933eb1b9c8ddb26512ead942ee2142676"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tests/resources/post_cell_execute_error.ipynb

@@ -0,0 +1,45 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "nbmake": {
+     "post_cell_execute": [
+      "raise Exception('boom!')"
+     ]
+    }
+   },
+   "outputs": [],
+   "source": [
+    "\"This cell has a post cell execution that fails\""
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.5"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "b62bdd1a52024cb952ae76a82719d36933eb1b9c8ddb26512ead942ee2142676"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tests/resources/post_cell_execute_masked_error.ipynb

@@ -0,0 +1,46 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "nbmake": {
+     "post_cell_execute": [
+      "raise Exception('boom!')"
+     ]
+    }
+   },
+   "outputs": [],
+   "source": [
+    "raise Exception(\"bang!\")\n",
+    "# this cell has a post_cell_execute that also raises an exception"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.5"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "b62bdd1a52024cb952ae76a82719d36933eb1b9c8ddb26512ead942ee2142676"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tests/test_nb_run.py

@@ -1,6 +1,6 @@
-import os
 from pathlib import Path
 
+import pytest
 from nbformat import write
 from nbformat.v4 import new_code_cell, new_notebook, new_output
 from pytest import Pytester
@@ -122,6 +122,33 @@ def test_when_mock_then_succeeds(self, testdir2: Never):
         res: NotebookResult = run.execute()
         assert res.error == None
 
+    def test_when_post_cell_execute_then_succeeds(self, testdir2: Never):
+        nb = Path(__file__).parent / "resources" / "post_cell_execute.ipynb"
+        run = NotebookRun(nb, 300)
+        res: NotebookResult = run.execute()
+        assert res.error == None
+
+    def test_when_post_cell_execute_then_command_fails(self, testdir2: Never):
+        nb = Path(__file__).parent / "resources" / "post_cell_execute_error.ipynb"
+        run = NotebookRun(nb, 300)
+        with pytest.raises(Exception) as exc_info:
+            run.execute()
+
+        assert exc_info != None
+        assert "boom!" in exc_info.value.args[0]
+
+    def test_when_post_cell_execute_then_cell_fails(self, testdir2: Never):
+        nb = (
+            Path(__file__).parent / "resources" / "post_cell_execute_masked_error.ipynb"
+        )
+        run = NotebookRun(nb, 300)
+        res: NotebookResult = run.execute()
+
+        # make sure the cell exception (bang!) is raised and not masked
+        # by the post cell execution exception (boom!)
+        assert res.error != None
+        assert "bang!" in res.error.summary
+
     def test_when_magic_error_then_fails(self, testdir2: Never):
         nb = Path(__file__).parent / "resources" / "magic_error.ipynb"
         run = NotebookRun(nb, 300)