remove cells emptied by nbqa (#744)

docs/history.rst

@@ -2,6 +2,10 @@
 Changelog
 =========
 
+1.5.0 (???)
+~~~~~~~~~~~
+``nbqa`` now removes empty cells which were not empty to begin with.
+
 1.4.0 (2022-07-17)
 ~~~~~~~~~~~~~~~~~~
 Added ``-nbqa-shell`` command (thanks @pwoolvett !).

nbqa/replace_source.py

@@ -196,17 +196,29 @@ def mutate(
     notebook_json = json.loads(notebook_txt)
     original_notebook_json = copy.deepcopy(notebook_json)
 
+    cells_to_remove = []
+
     cells = _get_cells(temp_file, len(notebook_info.temporary_lines), md=md)
     for code_cell_number, cell in enumerate(
         _notebook_cells(notebook_json, md=md), start=1
     ):
         if code_cell_number in notebook_info.code_cells_to_ignore:
             continue
-        cell["source"] = _get_new_source(code_cell_number, notebook_info, next(cells))
+        new_source = _get_new_source(code_cell_number, notebook_info, next(cells))
+        if not new_source:
+            cells_to_remove.append(code_cell_number)
+        cell["source"] = new_source
 
     if original_notebook_json == notebook_json:
         return False
 
+    if cells_to_remove:
+        notebook_json["cells"] = [
+            cell
+            for i, cell in enumerate(notebook_json["cells"], start=1)
+            if i not in cells_to_remove
+        ]
+
     temp_notebook = os.path.join(os.path.dirname(temp_file), os.path.basename(notebook))
     with open(temp_notebook, "w", encoding="utf-8") as handle:
         if notebook_txt.endswith("\n"):

tests/tools/test_isort_works.py

@@ -1,6 +1,6 @@
 """Check :code:`isort` works as intended."""
-
 import difflib
+import json
 import os
 from pathlib import Path
 from textwrap import dedent
@@ -242,3 +242,44 @@ def test_return_code_false_positive() -> None:
 
     result = main(["isort", str(notebook), "--nbqa-diff", "--float-to-top"])
     assert result == 1
+
+
+def test_float_to_top(tmp_test_data: Path) -> None:
+    """
+    Check isort works when a notebook has imports in different cells.
+
+    This test would fail if we didn't pass --treat-comment-as-code '# %%NBQA-CELL-SEP'.
+
+    Parameters
+    ----------
+    tmp_test_data
+        Temporary copy of test data.
+    """
+    notebook = tmp_test_data / "notebook_with_separated_imports_other.ipynb"
+
+    main(["isort", str(notebook), "--float-to-top"])
+
+    with open(notebook, encoding="utf-8") as fd:
+        result = json.load(fd)["cells"]
+    expected = [
+        {
+            "cell_type": "code",
+            "execution_count": None,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import os\n",
+                "\n",
+                "# This is a comment on the second import\n",
+                "import numpy",
+            ],
+        },
+        {
+            "cell_type": "code",
+            "execution_count": None,
+            "metadata": {},
+            "outputs": [],
+            "source": [],
+        },
+    ]
+    assert result == expected