remove correct cell in floattotop (#767)

Co-authored-by: MarcoGorelli <>

nbqa/replace_source.py

@@ -10,7 +10,17 @@
 import sys
 from difflib import unified_diff
 from shutil import move
-from typing import Any, Dict, Iterator, List, Mapping, MutableMapping, Sequence, Set
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Mapping,
+    MutableMapping,
+    Sequence,
+    Set,
+    Tuple,
+)
 
 import tokenize_rt
 
@@ -153,7 +163,7 @@ def _notebook_cells(
     notebook_json: Mapping[str, Any],
     *,
     md: bool,
-) -> Iterator[MutableMapping[str, Any]]:
+) -> Iterator[Tuple[int, MutableMapping[str, Any]]]:
     """
     Iterate through notebook's cells.
 
@@ -164,12 +174,14 @@ def _notebook_cells(
 
     Yields
     ------
+    cell_number
+        Index of cell (regardless of type).
     MutableMapping[str, Any]
         Cell content.
     """
-    for cell in notebook_json["cells"]:
+    for i, cell in enumerate(notebook_json["cells"]):
         if cell["cell_type"] == SOURCE[md]:
-            yield cell
+            yield i, cell
 
 
 def _write_notebook(
@@ -241,14 +253,14 @@ def mutate(
     cells_to_remove = []
 
     cells = _get_cells(temp_file, len(notebook_info.temporary_lines), md=md)
-    for code_cell_number, cell in enumerate(
+    for code_cell_number, (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
         new_source = _get_new_source(code_cell_number, notebook_info, next(cells))
         if not new_source:
-            cells_to_remove.append(code_cell_number)
+            cells_to_remove.append(cell_number)
         cell["source"] = new_source
 
     if original_notebook_json == notebook_json:
@@ -257,7 +269,7 @@ def mutate(
     if cells_to_remove:
         notebook_json["cells"] = [
             cell
-            for i, cell in enumerate(notebook_json["cells"], start=1)
+            for i, cell in enumerate(notebook_json["cells"])
             if i not in cells_to_remove
         ]
 
@@ -331,7 +343,7 @@ def diff(
 
     actually_mutated = False
 
-    for code_cell_number, cell in enumerate(
+    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:

tests/data/markdown_then_imports.ipynb

@@ -0,0 +1,39 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "hello world"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.6 ('.venv': venv)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.6"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "864e9cd7dbf32e150cc52d563f4c3a35053498b9e9d82e11d8edd2e505b61bc9"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tests/tools/test_isort_works.py

@@ -283,3 +283,33 @@ def test_float_to_top(tmp_test_data: Path) -> None:
         },
     ]
     assert result == expected
+
+
+def test_float_to_top_starting_markdown(tmp_test_data: Path) -> None:
+    """
+    Check isort works when a notebook has markdown in first cell.
+
+    The --float-to-top option would previously have removed the wrong cell.
+
+    Parameters
+    ----------
+    tmp_test_data
+        Temporary copy of test data.
+    """
+    notebook = tmp_test_data / "markdown_then_imports.ipynb"
+
+    main(["isort", str(notebook), "--float-to-top"])
+
+    with open(notebook, encoding="utf-8") as fd:
+        result = json.load(fd)["cells"]
+    expected = [
+        {"cell_type": "markdown", "metadata": {}, "source": ["hello world"]},
+        {
+            "cell_type": "code",
+            "execution_count": None,
+            "metadata": {},
+            "outputs": [],
+            "source": ["import os\n", "import sys"],
+        },
+    ]
+    assert result == expected