exec: write output notebook incrementally after each cell

src/nbexec/cli/exec_cmd.py

@@ -78,11 +78,12 @@ def _select_range(code_cells, total, from_cell, to_cell):
 # Cell execution loop
 # ---------------------------------------------------------------------------
 
-def _run_cells(session_id, selected, start, total, timeout):
+def _run_cells(session_id, selected, start, total, timeout, on_result=None):
     """Execute cells sequentially. Returns (results, interrupted).
 
     results is a list of (1-based cell index, result).
     Stops after the first cell that returns an error status or on Ctrl+C.
+    If on_result is provided, it is called with (cell_number, result) after each cell.
     """
     results = []
     interrupted = False
@@ -95,6 +96,8 @@ def _run_cells(session_id, selected, start, total, timeout):
             interrupted = True
             break
         results.append((i, result))
+        if on_result is not None:
+            on_result(i, result)
         if result.get("status") == "error":
             click.echo(f"Cell {i} failed, stopping.", err=True)
             break
@@ -186,13 +189,18 @@ def _exec_notebook(session_id, notebook_path, timeout, from_cell, to_cell, outpu
         click.echo(f"Executing {len(selected)} of {total} code cells from {notebook_path}", err=True)
         click.echo(f"Output notebook: {output_path}", err=True)
 
-    results, interrupted = _run_cells(session_id, selected, start, total, timeout)
-
+    # Load the output notebook upfront so we can write incrementally.
+    out_nb = None
     if output_path:
         out_nb = _load_output_base(output_path, notebook_path, is_partial)
-        _record_results(out_nb, results)
+
+    def _flush_result(cell_num, result):
+        _record_results(out_nb, [(cell_num, result)])
         _write_output_notebook(out_nb, output_path)
 
+    on_result = _flush_result if out_nb is not None else None
+    results, interrupted = _run_cells(session_id, selected, start, total, timeout, on_result)
+
     if interrupted:
         sys.exit(130)
     if results and results[-1][1].get("status") == "error":

tests/test_exec_notebook.py

@@ -513,6 +513,51 @@ def fake_send(method, params, timeout=None):
         assert len(code_cells[0].outputs) == 1  # first cell succeeded
         assert code_cells[1].outputs[0].output_type == "error"
 
+    def test_output_written_incrementally(self, tmp_path):
+        """Output notebook is flushed to disk after each cell, not just at the end."""
+        nb_path = tmp_path / "test.ipynb"
+        out_path = tmp_path / "out.ipynb"
+        _make_notebook(
+            [
+                new_code_cell(source="a = 1"),
+                new_code_cell(source="b = 2"),
+                new_code_cell(source="c = 3"),
+            ],
+            nb_path,
+        )
+
+        snapshots = []  # snapshot output notebook state after each cell
+
+        call_count = 0
+
+        def fake_send(method, params, timeout=None):
+            nonlocal call_count
+            call_count += 1
+            # After the first cell, the output file should already exist on disk.
+            if call_count > 1 and out_path.exists():
+                snapshots.append(_read_output_notebook(out_path))
+            return _make_daemon_response(
+                text=f"r{call_count}\n", execution_count=call_count,
+                outputs=[{"output_type": "stream", "name": "stdout", "text": f"r{call_count}\n"}],
+            )
+
+        runner = CliRunner()
+        with patch("nbexec.cli.exec_cmd.send_to_daemon", side_effect=fake_send):
+            result = runner.invoke(
+                exec_code,
+                ["--session", "s", "--file", str(nb_path), "--output", str(out_path)],
+            )
+
+        assert result.exit_code == 0
+        # We should have captured 2 snapshots (before cell 2 and before cell 3)
+        assert len(snapshots) == 2
+        # After cell 1 completed, snapshot should have cell 1 output
+        code_cells_snap1 = [c for c in snapshots[0].cells if c.cell_type == "code"]
+        assert len(code_cells_snap1[0].outputs) == 1
+        assert code_cells_snap1[0].outputs[0].text == "r1\n"
+        # Cell 2 hasn't been recorded yet in this snapshot
+        assert len(code_cells_snap1[1].outputs) == 0
+
     def test_output_rejected_for_non_notebook(self, tmp_path):
         """--output only allowed with .ipynb files."""
         py_path = tmp_path / "script.py"