diff --git a/src/imports/ipynb.rs b/src/imports/ipynb.rs index 596278bd..5bed38af 100644 --- 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 index a51bdb9d..47ebeb1b 100644 --- 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 index 28f3fbc0..55e0fef9 100644 --- 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 index 5e9dc73d..0c327602 100644 --- 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)], }