fix(imports): remove invalid lines when parsing notebooks (#656)

fpgmaas · Apr 2, 2024 · a40f47b · a40f47b
1 parent 5494d09
commit a40f47b
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 5 deletions.
diff --git a/src/imports/ipynb.rs b/src/imports/ipynb.rs
@@ -66,6 +66,19 @@ fn _extract_code_from_notebook_cells(cells: &[serde_json::Value]) -> String {
         .filter_map(|cell| cell["source"].as_array())
         .flatten()
         .filter_map(|line| line.as_str())
+        .map(|line| {
+            // https://ipython.readthedocs.io/en/stable/interactive/magics.html
+            // We want to skip lines using magics, as they are not valid Python. We replace the
+            // lines with empty strings instead of completely removing them, to ensure that the
+            // violation reporters show the correct line location for imports made below those
+            // lines.
+            if line.starts_with('%') || line.starts_with('!') {
+                ""
+            } else {
+                line
+            }
+        })
+        .map(|s| s.strip_suffix('\n').unwrap_or(s))
         .map(str::to_owned)
         .collect();
 

diff --git a/tests/data/example_project/src/notebook.ipynb b/tests/data/example_project/src/notebook.ipynb
@@ -7,9 +7,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "!ls\n",
+    "%timeit\n",
+    "%%timeit\n",
     "import click\n",
     "from urllib3 import contrib\n",
-    "import toml"
+    "import toml\n",
+    "1 +\\\n",
+    "    2"
    ]
   }
  ],

diff --git a/tests/functional/cli/test_cli_requirements_txt.py b/tests/functional/cli/test_cli_requirements_txt.py
@@ -97,7 +97,7 @@ def test_cli_single_requirements_files(pip_venv_factory: PipVenvFactory) -> None
                 "module": "urllib3",
                 "location": {
                     "file": str(Path("src/notebook.ipynb")),
-                    "line": 3,
+                    "line": 2,
                     "column": 1,
                 },
             },

diff --git a/tests/unit/imports/test_extract.py b/tests/unit/imports/test_extract.py
@@ -49,9 +49,9 @@ def test_import_parser_ipynb() -> None:
     notebook_path = Path("tests/data/example_project/src/notebook.ipynb")
 
     assert get_imported_modules_from_list_of_files([notebook_path]) == {
-        "click": [Location(notebook_path, 1, 8)],
-        "toml": [Location(notebook_path, 5, 8)],
-        "urllib3": [Location(notebook_path, 3, 1)],
+        "click": [Location(notebook_path, 4, 8)],
+        "toml": [Location(notebook_path, 6, 8)],
+        "urllib3": [Location(notebook_path, 5, 1)],
     }