Update to ensure no hanging code cells in generated notebooks

jupyterlab · Dec 18, 2024 · 9a48bd9 · 9a48bd9
1 parent 1cbd853
commit 9a48bd9
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/generate.py
@@ -212,6 +212,15 @@ def create_notebook(outline):
         nb["cells"].append(nbf.new_markdown_cell("## " + section["title"]))
         for code_block in section["code"].split("\n\n"):
             nb["cells"].append(nbf.new_code_cell(code_block))
+
+    # Post process notebook for hanging cells: merge hanging cell with the previous cell
+    nb_cells = list()
+    for cell in nb["cells"]:
+        if (cell["cell_type"] == 'code') and (cell["source"][0] == ' '):
+            nb_cells[-1]["source"] = nb_cells[-1]["source"] + '\n\n' + cell["source"]
+        else:
+            nb_cells.append(cell)
+    nb["cells"] = nb_cells
     return nb