Create ruff_notebook crate

CONTRIBUTING.md

@@ -129,6 +129,7 @@ At time of writing, the repository includes the following crates:
     intermediate representation. The backend for `ruff_python_formatter`.
 - `crates/ruff_index`: library crate inspired by `rustc_index`.
 - `crates/ruff_macros`: proc macro crate containing macros used by Ruff.
+- `crates/ruff_notebook`: library crate for parsing and manipulating Jupyter notebooks.
 - `crates/ruff_python_ast`: library crate containing Python-specific AST types and utilities.
 - `crates/ruff_python_codegen`: library crate containing utilities for generating Python source code.
 - `crates/ruff_python_formatter`: library crate implementing the Python formatter. Emits an

crates/ruff/Cargo.toml

@@ -18,6 +18,7 @@ name = "ruff"
 ruff_cache = { path = "../ruff_cache" }
 ruff_diagnostics = { path = "../ruff_diagnostics", features = ["serde"] }
 ruff_index = { path = "../ruff_index" }
+ruff_notebook = { path = "../ruff_notebook" }
 ruff_macros = { path = "../ruff_macros" }
 ruff_python_ast = { path = "../ruff_python_ast", features = ["serde"] }
 ruff_python_codegen = { path = "../ruff_python_codegen" }
@@ -64,17 +65,15 @@ schemars = { workspace = true, optional = true }
 semver = { version = "1.0.16" }
 serde = { workspace = true }
 serde_json = { workspace = true }
-serde_with = { version = "3.0.0" }
 similar = { workspace = true }
 smallvec = { workspace = true }
 strum = { workspace = true }
 strum_macros = { workspace = true }
-thiserror = { version = "1.0.43" }
+thiserror = { workspace = true }
 toml = { workspace = true }
 typed-arena = { version = "2.0.2" }
 unicode-width = { workspace = true }
 unicode_names2 = { version = "0.6.0", git = "https://github.com/youknowone/unicode_names2.git", rev = "4ce16aa85cbcdd9cc830410f1a72ef9a235f2fde" }
-uuid = { workspace = true, features = ["v4", "fast-rng", "macro-diagnostics", "js"] }
 wsl = { version = "0.1.0" }
 
 [dev-dependencies]

crates/ruff/src/autofix/mod.rs

@@ -4,17 +4,15 @@ use std::collections::BTreeSet;
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 use rustc_hash::{FxHashMap, FxHashSet};
 
-use ruff_diagnostics::{Diagnostic, Edit, Fix, IsolationLevel};
+use ruff_diagnostics::{Diagnostic, Edit, Fix, IsolationLevel, SourceMap};
 use ruff_source_file::Locator;
 
-use crate::autofix::source_map::SourceMap;
 use crate::linter::FixTable;
 use crate::registry::{AsRule, Rule};
 
 pub(crate) mod codemods;
 pub(crate) mod edits;
 pub(crate) mod snippet;
-pub(crate) mod source_map;
 
 pub(crate) struct FixResult {
     /// The resulting source code, after applying all fixes.
@@ -140,10 +138,9 @@ fn cmp_fix(rule1: Rule, rule2: Rule, fix1: &Fix, fix2: &Fix) -> std::cmp::Orderi
 mod tests {
     use ruff_text_size::{Ranged, TextSize};
 
-    use ruff_diagnostics::{Diagnostic, Edit, Fix};
+    use ruff_diagnostics::{Diagnostic, Edit, Fix, SourceMarker};
     use ruff_source_file::Locator;
 
-    use crate::autofix::source_map::SourceMarker;
     use crate::autofix::{apply_fixes, FixResult};
     use crate::rules::pycodestyle::rules::MissingNewlineAtEndOfFile;
 
@@ -207,14 +204,8 @@ print("hello world")
         assert_eq!(
             source_map.markers(),
             &[
-                SourceMarker {
-                    source: 10.into(),
-                    dest: 10.into(),
-                },
-                SourceMarker {
-                    source: 10.into(),
-                    dest: 21.into(),
-                },
+                SourceMarker::new(10.into(), 10.into(),),
+                SourceMarker::new(10.into(), 21.into(),),
             ]
         );
     }
@@ -250,14 +241,8 @@ class A(Bar):
         assert_eq!(
             source_map.markers(),
             &[
-                SourceMarker {
-                    source: 8.into(),
-                    dest: 8.into(),
-                },
-                SourceMarker {
-                    source: 14.into(),
-                    dest: 11.into(),
-                },
+                SourceMarker::new(8.into(), 8.into(),),
+                SourceMarker::new(14.into(), 11.into(),),
             ]
         );
     }
@@ -289,14 +274,8 @@ class A:
         assert_eq!(
             source_map.markers(),
             &[
-                SourceMarker {
-                    source: 7.into(),
-                    dest: 7.into()
-                },
-                SourceMarker {
-                    source: 15.into(),
-                    dest: 7.into()
-                }
+                SourceMarker::new(7.into(), 7.into()),
+                SourceMarker::new(15.into(), 7.into()),
             ]
         );
     }
@@ -332,22 +311,10 @@ class A(object):
         assert_eq!(
             source_map.markers(),
             &[
-                SourceMarker {
-                    source: 8.into(),
-                    dest: 8.into()
-                },
-                SourceMarker {
-                    source: 16.into(),
-                    dest: 8.into()
-                },
-                SourceMarker {
-                    source: 22.into(),
-                    dest: 14.into(),
-                },
-                SourceMarker {
-                    source: 30.into(),
-                    dest: 14.into(),
-                }
+                SourceMarker::new(8.into(), 8.into()),
+                SourceMarker::new(16.into(), 8.into()),
+                SourceMarker::new(22.into(), 14.into(),),
+                SourceMarker::new(30.into(), 14.into(),),
             ]
         );
     }
@@ -382,14 +349,8 @@ class A:
         assert_eq!(
             source_map.markers(),
             &[
-                SourceMarker {
-                    source: 7.into(),
-                    dest: 7.into(),
-                },
-                SourceMarker {
-                    source: 15.into(),
-                    dest: 7.into(),
-                }
+                SourceMarker::new(7.into(), 7.into(),),
+                SourceMarker::new(15.into(), 7.into(),),
             ]
         );
     }

crates/ruff/src/lib.rs

@@ -20,7 +20,6 @@ mod doc_lines;
 mod docstrings;
 pub mod fs;
 mod importer;
-pub mod jupyter;
 mod lex;
 pub mod line_width;
 pub mod linter;

crates/ruff/src/linter.rs

@@ -6,16 +6,15 @@ use anyhow::{anyhow, Result};
 use colored::Colorize;
 use itertools::Itertools;
 use log::error;
-use ruff_python_parser::lexer::LexResult;
-use ruff_python_parser::{AsMode, ParseError};
 use rustc_hash::FxHashMap;
 
 use ruff_diagnostics::Diagnostic;
 use ruff_python_ast::imports::ImportMap;
 use ruff_python_ast::PySourceType;
 use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
-
+use ruff_python_parser::lexer::LexResult;
+use ruff_python_parser::{AsMode, ParseError};
 use ruff_source_file::{Locator, SourceFileBuilder};
 use ruff_text_size::Ranged;
 
@@ -609,3 +608,133 @@ This indicates a bug in `{}`. If you could open an issue at:
         );
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::path::Path;
+
+    use anyhow::Result;
+    use test_case::test_case;
+
+    use ruff_notebook::{Notebook, NotebookError};
+
+    use crate::registry::Rule;
+    use crate::source_kind::SourceKind;
+    use crate::test::{test_contents, test_notebook_path, TestedNotebook};
+    use crate::{assert_messages, settings};
+
+    /// Construct a path to a Jupyter notebook in the `resources/test/fixtures/jupyter` directory.
+    fn notebook_path(path: impl AsRef<Path>) -> std::path::PathBuf {
+        Path::new("../ruff_notebook/resources/test/fixtures/jupyter").join(path)
+    }
+
+    #[test]
+    fn test_import_sorting() -> Result<(), NotebookError> {
+        let actual = notebook_path("isort.ipynb");
+        let expected = notebook_path("isort_expected.ipynb");
+        let TestedNotebook {
+            messages,
+            source_notebook,
+            ..
+        } = test_notebook_path(
+            &actual,
+            expected,
+            &settings::Settings::for_rule(Rule::UnsortedImports),
+        )?;
+        assert_messages!(messages, actual, source_notebook);
+        Ok(())
+    }
+
+    #[test]
+    fn test_ipy_escape_command() -> Result<(), NotebookError> {
+        let actual = notebook_path("ipy_escape_command.ipynb");
+        let expected = notebook_path("ipy_escape_command_expected.ipynb");
+        let TestedNotebook {
+            messages,
+            source_notebook,
+            ..
+        } = test_notebook_path(
+            &actual,
+            expected,
+            &settings::Settings::for_rule(Rule::UnusedImport),
+        )?;
+        assert_messages!(messages, actual, source_notebook);
+        Ok(())
+    }
+
+    #[test]
+    fn test_unused_variable() -> Result<(), NotebookError> {
+        let actual = notebook_path("unused_variable.ipynb");
+        let expected = notebook_path("unused_variable_expected.ipynb");
+        let TestedNotebook {
+            messages,
+            source_notebook,
+            ..
+        } = test_notebook_path(
+            &actual,
+            expected,
+            &settings::Settings::for_rule(Rule::UnusedVariable),
+        )?;
+        assert_messages!(messages, actual, source_notebook);
+        Ok(())
+    }
+
+    #[test]
+    fn test_json_consistency() -> Result<()> {
+        let actual_path = notebook_path("before_fix.ipynb");
+        let expected_path = notebook_path("after_fix.ipynb");
+
+        let TestedNotebook {
+            linted_notebook: fixed_notebook,
+            ..
+        } = test_notebook_path(
+            actual_path,
+            &expected_path,
+            &settings::Settings::for_rule(Rule::UnusedImport),
+        )?;
+        let mut writer = Vec::new();
+        fixed_notebook.write(&mut writer)?;
+        let actual = String::from_utf8(writer)?;
+        let expected = std::fs::read_to_string(expected_path)?;
+        assert_eq!(actual, expected);
+        Ok(())
+    }
+
+    #[test_case(Path::new("before_fix.ipynb"), true; "trailing_newline")]
+    #[test_case(Path::new("no_trailing_newline.ipynb"), false; "no_trailing_newline")]
+    fn test_trailing_newline(path: &Path, trailing_newline: bool) -> Result<()> {
+        let notebook = Notebook::from_path(&notebook_path(path))?;
+        assert_eq!(notebook.trailing_newline(), trailing_newline);
+
+        let mut writer = Vec::new();
+        notebook.write(&mut writer)?;
+        let string = String::from_utf8(writer)?;
+        assert_eq!(string.ends_with('\n'), trailing_newline);
+
+        Ok(())
+    }
+
+    // Version <4.5, don't emit cell ids
+    #[test_case(Path::new("no_cell_id.ipynb"), false; "no_cell_id")]
+    // Version 4.5, cell ids are missing and need to be added
+    #[test_case(Path::new("add_missing_cell_id.ipynb"), true; "add_missing_cell_id")]
+    fn test_cell_id(path: &Path, has_id: bool) -> Result<()> {
+        let source_notebook = Notebook::from_path(&notebook_path(path))?;
+        let source_kind = SourceKind::IpyNotebook(source_notebook);
+        let (_, transformed) = test_contents(
+            &source_kind,
+            path,
+            &settings::Settings::for_rule(Rule::UnusedImport),
+        );
+        let linted_notebook = transformed.into_owned().expect_ipy_notebook();
+        let mut writer = Vec::new();
+        linted_notebook.write(&mut writer)?;
+        let actual = String::from_utf8(writer)?;
+        if has_id {
+            assert!(actual.contains(r#""id": ""#));
+        } else {
+            assert!(!actual.contains(r#""id":"#));
+        }
+        Ok(())
+    }
+}

crates/ruff/src/logging.rs

@@ -12,8 +12,8 @@ use ruff_python_parser::{ParseError, ParseErrorType};
 use ruff_source_file::{OneIndexed, SourceCode, SourceLocation};
 
 use crate::fs;
-use crate::jupyter::Notebook;
 use crate::source_kind::SourceKind;
+use ruff_notebook::Notebook;
 
 pub static WARNINGS: Lazy<Mutex<Vec<&'static str>>> = Lazy::new(Mutex::default);