diff --git a/crates/ruff/src/jupyter/notebook.rs b/crates/ruff/src/jupyter/notebook.rs index 66c4240fe0..70420e89a7 100644 --- a/crates/ruff/src/jupyter/notebook.rs +++ b/crates/ruff/src/jupyter/notebook.rs @@ -1,6 +1,6 @@ use std::cmp::Ordering; use std::fs::File; -use std::io::{BufReader, BufWriter}; +use std::io::{BufReader, BufWriter, Cursor, Write}; use std::iter; use std::path::Path; @@ -23,6 +23,22 @@ 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 { + let mut notebook = Notebook::read(path).map_err(|err| { + anyhow::anyhow!( + "Failed to read notebook file `{}`: {:?}", + path.display(), + err + ) + })?; + let code = notebook.content().to_string(); + notebook.update_cell_content(&code); + let mut buffer = Cursor::new(Vec::new()); + notebook.write_inner(&mut buffer)?; + Ok(String::from_utf8(buffer.into_inner())?) +} + /// Return `true` if the [`Path`] appears to be that of a jupyter notebook file (`.ipynb`). pub fn is_jupyter_notebook(path: &Path) -> bool { path.extension() @@ -370,13 +386,18 @@ impl Notebook { .map_or(true, |language| language.name == "python") } + fn write_inner(&self, writer: &mut impl Write) -> anyhow::Result<()> { + // https://github.com/psf/black/blob/69ca0a4c7a365c5f5eea519a90980bab72cab764/src/black/__init__.py#LL1041 + let formatter = serde_json::ser::PrettyFormatter::with_indent(b" "); + let mut ser = serde_json::Serializer::with_formatter(writer, formatter); + self.raw.serialize(&mut ser)?; + Ok(()) + } + /// Write back with an indent of 1, just like black pub fn write(&self, path: &Path) -> anyhow::Result<()> { let mut writer = BufWriter::new(File::create(path)?); - // https://github.com/psf/black/blob/69ca0a4c7a365c5f5eea519a90980bab72cab764/src/black/__init__.py#LL1041 - let formatter = serde_json::ser::PrettyFormatter::with_indent(b" "); - let mut ser = serde_json::Serializer::with_formatter(&mut writer, formatter); - self.raw.serialize(&mut ser)?; + self.write_inner(&mut writer)?; Ok(()) } } diff --git a/crates/ruff_dev/src/round_trip.rs b/crates/ruff_dev/src/round_trip.rs index 4cc4f36c59..75d1cf59e4 100644 --- a/crates/ruff_dev/src/round_trip.rs +++ b/crates/ruff_dev/src/round_trip.rs @@ -1,4 +1,4 @@ -//! Run round-trip source code generation on a given Python file. +//! Run round-trip source code generation on a given Python or Jupyter notebook file. #![allow(clippy::print_stdout, clippy::print_stderr)] use std::fs; @@ -6,17 +6,23 @@ use std::path::PathBuf; use anyhow::Result; +use ruff::jupyter; use ruff::round_trip; #[derive(clap::Args)] pub(crate) struct Args { - /// Python file to round-trip. + /// Python or Jupyter notebook file to round-trip. #[arg(required = true)] file: PathBuf, } pub(crate) fn main(args: &Args) -> Result<()> { - let contents = fs::read_to_string(&args.file)?; - println!("{}", round_trip(&contents, &args.file.to_string_lossy())?); + let path = args.file.as_path(); + if jupyter::is_jupyter_notebook(path) { + println!("{}", jupyter::round_trip(path)?); + } else { + let contents = fs::read_to_string(&args.file)?; + println!("{}", round_trip(&contents, &args.file.to_string_lossy())?); + } Ok(()) }