Use Jupyter mode for the parser with Notebook files

crates/ruff/resources/test/fixtures/jupyter/cell/cell_magic.json

@@ -0,0 +1,8 @@
+{
+  "execution_count": null,
+  "cell_type": "code",
+  "id": "1",
+  "metadata": {},
+  "outputs": [],
+  "source": ["%%timeit\n", "print('hello world')"]
+}

crates/ruff/resources/test/fixtures/jupyter/line_magics.ipynb

@@ -0,0 +1,41 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eab4754a-d6df-4b41-8ee8-7e23aef440f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "\n",
+    "%matplotlib inline\n",
+    "\n",
+    "import os\n",
+    "\n",
+    "_ = math.pi"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (ruff)",
+   "language": "python",
+   "name": "ruff"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

crates/ruff/resources/test/fixtures/jupyter/line_magics_expected.ipynb

@@ -0,0 +1,40 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cad32845-44f9-4a53-8b8c-a6b1bb3f3378",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "\n",
+    "%matplotlib inline\n",
+    "\n",
+    "\n",
+    "_ = math.pi"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python (ruff)",
+   "language": "python",
+   "name": "ruff"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

crates/ruff/src/importer/insertion.rs

@@ -302,7 +302,7 @@ mod tests {
     use ruff_text_size::TextSize;
     use rustpython_parser::ast::Suite;
     use rustpython_parser::lexer::LexResult;
-    use rustpython_parser::Parse;
+    use rustpython_parser::{Mode, Parse};
 
     use ruff_python_ast::source_code::{Locator, Stylist};
     use ruff_python_trivia::LineEnding;
@@ -313,7 +313,7 @@ mod tests {
     fn start_of_file() -> Result<()> {
         fn insert(contents: &str) -> Result<Insertion> {
             let program = Suite::parse(contents, "<filename>")?;
-            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents);
+            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             Ok(Insertion::start_of_file(&program, &locator, &stylist))
@@ -424,7 +424,7 @@ x = 1
     #[test]
     fn start_of_block() {
         fn insert(contents: &str, offset: TextSize) -> Insertion {
-            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents);
+            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents, Mode::Module);
             let locator = Locator::new(contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             Insertion::start_of_block(offset, &locator, &stylist)

crates/ruff/src/jupyter/notebook.rs

@@ -6,6 +6,7 @@ use std::path::Path;
 
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
+use rustpython_parser::Mode;
 use serde::Serialize;
 use serde_json::error::Category;
 
@@ -21,8 +22,6 @@ use crate::IOError;
 
 pub const JUPYTER_NOTEBOOK_EXT: &str = "ipynb";
 
-const MAGIC_PREFIX: [&str; 3] = ["%", "!", "?"];
-
 /// Run round-trip source code generation on a given Jupyter notebook file path.
 pub fn round_trip(path: &Path) -> anyhow::Result<String> {
     let mut notebook = Notebook::read(path).map_err(|err| {
@@ -61,26 +60,21 @@ impl Cell {
     /// Return `true` if it's a valid code cell.
     ///
     /// A valid code cell is a cell where the cell type is [`Cell::Code`] and the
-    /// source doesn't contain a magic, shell or help command.
+    /// source doesn't contain a cell magic.
     fn is_valid_code_cell(&self) -> bool {
         let source = match self {
             Cell::Code(cell) => &cell.source,
             _ => return false,
         };
-        // Ignore a cell if it contains a magic command. There could be valid
-        // Python code as well, but we'll ignore that for now.
-        // TODO(dhruvmanila): https://github.com/psf/black/blob/main/src/black/handle_ipynb_magics.py
+        // Ignore cells containing cell magic. This is different from line magic
+        // which is allowed and ignored by the parser.
         !match source {
-            SourceValue::String(string) => string.lines().any(|line| {
-                MAGIC_PREFIX
-                    .iter()
-                    .any(|prefix| line.trim_start().starts_with(prefix))
-            }),
-            SourceValue::StringArray(string_array) => string_array.iter().any(|line| {
-                MAGIC_PREFIX
-                    .iter()
-                    .any(|prefix| line.trim_start().starts_with(prefix))
-            }),
+            SourceValue::String(string) => string
+                .lines()
+                .any(|line| line.trim_start().starts_with("%%")),
+            SourceValue::StringArray(string_array) => string_array
+                .iter()
+                .any(|line| line.trim_start().starts_with("%%")),
         }
     }
 }
@@ -158,7 +152,7 @@ impl Notebook {
                                 )
                             })?;
                             // Check if tokenizing was successful and the file is non-empty
-                            if (ruff_rustpython::tokenize(&contents))
+                            if (ruff_rustpython::tokenize(&contents, Mode::Module))
                                 .last()
                                 .map_or(true, Result::is_err)
                             {
@@ -494,9 +488,10 @@ mod tests {
     }
 
     #[test_case(Path::new("markdown.json"), false; "markdown")]
-    #[test_case(Path::new("only_magic.json"), false; "only_magic")]
-    #[test_case(Path::new("code_and_magic.json"), false; "code_and_magic")]
+    #[test_case(Path::new("only_magic.json"), true; "only_magic")]
+    #[test_case(Path::new("code_and_magic.json"), true; "code_and_magic")]
     #[test_case(Path::new("only_code.json"), true; "only_code")]
+    #[test_case(Path::new("cell_magic.json"), false; "cell_magic")]
     fn test_is_valid_code_cell(path: &Path, expected: bool) -> Result<()> {
         assert_eq!(read_jupyter_cell(path)?.is_valid_code_cell(), expected);
         Ok(())
@@ -557,6 +552,18 @@ print("after empty cells")
         Ok(())
     }
 
+    #[test]
+    fn test_line_magics() -> Result<()> {
+        let path = "line_magics.ipynb".to_string();
+        let (diagnostics, source_kind) = test_notebook_path(
+            &path,
+            Path::new("line_magics_expected.ipynb"),
+            &settings::Settings::for_rule(Rule::UnusedImport),
+        )?;
+        assert_messages!(diagnostics, path, source_kind);
+        Ok(())
+    }
+
     #[test]
     fn test_json_consistency() -> Result<()> {
         let path = "before_fix.ipynb".to_string();

crates/ruff/src/jupyter/snapshots/ruff__jupyter__notebook__tests__line_magics.snap

@@ -0,0 +1,21 @@
+---
+source: crates/ruff/src/jupyter/notebook.rs
+---
+line_magics.ipynb:cell 1:5:8: F401 [*] `os` imported but unused
+  |
+3 | %matplotlib inline
+4 | 
+5 | import os
+  |        ^^ F401
+  |
+  = help: Remove unused import: `os`
+
+ℹ Fix
+2 2 | 
+3 3 | %matplotlib inline
+4 4 | 
+5   |-import os
+6 5 | 
+7 6 | _ = math.pi
+
+

crates/ruff/src/linter.rs

@@ -8,7 +8,7 @@ use itertools::Itertools;
 use log::error;
 use rustc_hash::FxHashMap;
 use rustpython_parser::lexer::LexResult;
-use rustpython_parser::ParseError;
+use rustpython_parser::{Mode, ParseError};
 
 use ruff_diagnostics::Diagnostic;
 use ruff_python_ast::imports::ImportMap;
@@ -136,7 +136,11 @@ pub fn check_path(
             .iter_enabled()
             .any(|rule_code| rule_code.lint_source().is_imports());
     if use_ast || use_imports || use_doc_lines {
-        match ruff_rustpython::parse_program_tokens(tokens, &path.to_string_lossy()) {
+        match ruff_rustpython::parse_program_tokens(
+            tokens,
+            &path.to_string_lossy(),
+            source_kind.map_or(false, SourceKind::is_jupyter),
+        ) {
             Ok(python_ast) => {
                 if use_ast {
                     diagnostics.extend(check_ast(
@@ -258,7 +262,7 @@ pub fn add_noqa_to_path(path: &Path, package: Option<&Path>, settings: &Settings
     let contents = std::fs::read_to_string(path)?;
 
     // Tokenize once.
-    let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&contents);
+    let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&contents, Mode::Module);
 
     // Map row and column locations to byte slices (lazily).
     let locator = Locator::new(&contents);
@@ -325,8 +329,14 @@ pub fn lint_only(
     noqa: flags::Noqa,
     source_kind: Option<&SourceKind>,
 ) -> LinterResult<(Vec<Message>, Option<ImportMap>)> {
+    let mode = if source_kind.map_or(false, SourceKind::is_jupyter) {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+
     // Tokenize once.
-    let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents);
+    let tokens: Vec<LexResult> = ruff_rustpython::tokenize(contents, mode);
 
     // Map row and column locations to byte slices (lazily).
     let locator = Locator::new(contents);
@@ -415,10 +425,16 @@ pub fn lint_fix<'a>(
     // Track whether the _initial_ source code was parseable.
     let mut parseable = false;
 
+    let mode = if source_kind.is_jupyter() {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+
     // Continuously autofix until the source code stabilizes.
     loop {
         // Tokenize once.
-        let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&transformed);
+        let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&transformed, mode);
 
         // Map row and column locations to byte slices (lazily).
         let locator = Locator::new(&transformed);

crates/ruff/src/rules/pyflakes/mod.rs

@@ -12,6 +12,7 @@ mod tests {
     use anyhow::Result;
     use regex::Regex;
     use rustpython_parser::lexer::LexResult;
+    use rustpython_parser::Mode;
     use test_case::test_case;
 
     use ruff_diagnostics::Diagnostic;
@@ -502,7 +503,7 @@ mod tests {
     fn flakes(contents: &str, expected: &[Rule]) {
         let contents = dedent(contents);
         let settings = Settings::for_rules(Linter::Pyflakes.rules());
-        let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&contents);
+        let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&contents, Mode::Module);
         let locator = Locator::new(&contents);
         let stylist = Stylist::from_tokens(&tokens, &locator);
         let indexer = Indexer::from_tokens(&tokens, &locator);

crates/ruff/src/test.rs

@@ -9,6 +9,7 @@ use itertools::Itertools;
 use ruff_textwrap::dedent;
 use rustc_hash::FxHashMap;
 use rustpython_parser::lexer::LexResult;
+use rustpython_parser::Mode;
 
 use ruff_diagnostics::{AutofixKind, Diagnostic};
 use ruff_python_ast::source_code::{Indexer, Locator, SourceFileBuilder, Stylist};
@@ -97,8 +98,13 @@ pub(crate) fn max_iterations() -> usize {
 /// A convenient wrapper around [`check_path`], that additionally
 /// asserts that autofixes converge after a fixed number of iterations.
 fn test_contents(source_kind: &mut SourceKind, path: &Path, settings: &Settings) -> Vec<Message> {
+    let mode = if source_kind.is_jupyter() {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
     let contents = source_kind.content().to_string();
-    let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&contents);
+    let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&contents, mode);
     let locator = Locator::new(&contents);
     let stylist = Stylist::from_tokens(&tokens, &locator);
     let indexer = Indexer::from_tokens(&tokens, &locator);
@@ -160,7 +166,7 @@ fn test_contents(source_kind: &mut SourceKind, path: &Path, settings: &Settings)
                 notebook.update(&source_map, &fixed_contents);
             };
 
-            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&fixed_contents);
+            let tokens: Vec<LexResult> = ruff_rustpython::tokenize(&fixed_contents, mode);
             let locator = Locator::new(&fixed_contents);
             let stylist = Stylist::from_tokens(&tokens, &locator);
             let indexer = Indexer::from_tokens(&tokens, &locator);

crates/ruff_dev/src/print_ast.rs

@@ -5,19 +5,26 @@ use std::fs;
 use std::path::PathBuf;
 
 use anyhow::Result;
-use rustpython_parser::ast::Suite;
-use rustpython_parser::Parse;
+use rustpython_parser::{parse, Mode};
 
 #[derive(clap::Args)]
 pub(crate) struct Args {
     /// Python file for which to generate the AST.
     #[arg(required = true)]
     file: PathBuf,
+    /// Run in Jupyter mode i.e., allow line magics.
+    #[arg(long)]
+    jupyter: bool,
 }
 
 pub(crate) fn main(args: &Args) -> Result<()> {
     let contents = fs::read_to_string(&args.file)?;
-    let python_ast = Suite::parse(&contents, &args.file.to_string_lossy())?;
+    let mode = if args.jupyter {
+        Mode::Jupyter
+    } else {
+        Mode::Module
+    };
+    let python_ast = parse(&contents, mode, &args.file.to_string_lossy())?;
     println!("{python_ast:#?}");
     Ok(())
 }