From 56d445add9ade4d5ef56ec71618d32fe4e54768c Mon Sep 17 00:00:00 2001 From: Meheret <101792782+senadev42@users.noreply.github.com> Date: Fri, 1 Mar 2024 11:55:30 +0300 Subject: [PATCH] Colorize the output of ruff format --diff (#10110) Co-authored-by: Micha Reiser --- crates/ruff/src/commands/format.rs | 2 +- crates/ruff/src/commands/format_stdin.rs | 10 +- crates/ruff/src/diagnostics.rs | 13 +- crates/ruff_linter/src/source_kind.rs | 167 +++++++++++++++++------ 4 files changed, 146 insertions(+), 46 deletions(-) diff --git a/crates/ruff/src/commands/format.rs b/crates/ruff/src/commands/format.rs index f760ec96a14a5..185d79ff1f909 100644 --- a/crates/ruff/src/commands/format.rs +++ b/crates/ruff/src/commands/format.rs @@ -532,7 +532,7 @@ impl<'a> FormatResults<'a> { }) .sorted_unstable_by_key(|(path, _, _)| *path) { - unformatted.diff(formatted, Some(path), f)?; + write!(f, "{}", unformatted.diff(formatted, Some(path)).unwrap())?; } Ok(()) diff --git a/crates/ruff/src/commands/format_stdin.rs b/crates/ruff/src/commands/format_stdin.rs index f23b459c3aae8..bb3f7eb7bebd6 100644 --- a/crates/ruff/src/commands/format_stdin.rs +++ b/crates/ruff/src/commands/format_stdin.rs @@ -122,9 +122,13 @@ fn format_source_code( } FormatMode::Check => {} FormatMode::Diff => { - source_kind - .diff(formatted, path, &mut stdout().lock()) - .map_err(|err| FormatCommandError::Diff(path.map(Path::to_path_buf), err))?; + use std::io::Write; + write!( + &mut stdout().lock(), + "{}", + source_kind.diff(formatted, path).unwrap() + ) + .map_err(|err| FormatCommandError::Diff(path.map(Path::to_path_buf), err))?; } }, FormattedSource::Unchanged => { diff --git a/crates/ruff/src/diagnostics.rs b/crates/ruff/src/diagnostics.rs index 58d26e894ad14..ad15e88a30bcc 100644 --- a/crates/ruff/src/diagnostics.rs +++ b/crates/ruff/src/diagnostics.rs @@ -3,6 +3,7 @@ use std::borrow::Cow; use std::fs::File; use std::io; +use std::io::Write; use std::ops::{Add, AddAssign}; use std::path::Path; @@ -289,10 +290,10 @@ pub(crate) fn lint_path( match fix_mode { flags::FixMode::Apply => transformed.write(&mut File::create(path)?)?, flags::FixMode::Diff => { - source_kind.diff( - transformed.as_ref(), - Some(path), + write!( &mut io::stdout().lock(), + "{}", + source_kind.diff(&transformed, Some(path)).unwrap() )?; } flags::FixMode::Generate => {} @@ -442,7 +443,11 @@ pub(crate) fn lint_stdin( flags::FixMode::Diff => { // But only write a diff if it's non-empty. if !fixed.is_empty() { - source_kind.diff(transformed.as_ref(), path, &mut io::stdout().lock())?; + write!( + &mut io::stdout().lock(), + "{}", + source_kind.diff(&transformed, path).unwrap() + )?; } } flags::FixMode::Generate => {} diff --git a/crates/ruff_linter/src/source_kind.rs b/crates/ruff_linter/src/source_kind.rs index 892eb5b595467..c59f6eb224330 100644 --- a/crates/ruff_linter/src/source_kind.rs +++ b/crates/ruff_linter/src/source_kind.rs @@ -1,15 +1,18 @@ +use std::fmt::Formatter; use std::io; use std::io::Write; use std::path::Path; use anyhow::Result; -use similar::TextDiff; +use similar::{ChangeTag, TextDiff}; use thiserror::Error; use ruff_diagnostics::SourceMap; use ruff_notebook::{Cell, Notebook, NotebookError}; use ruff_python_ast::PySourceType; +use colored::Colorize; + use crate::fs; #[derive(Clone, Debug, PartialEq, is_macro::Is)] @@ -87,33 +90,53 @@ impl SourceKind { } } - /// Write a diff of the transformed source file to `stdout`. - pub fn diff( - &self, - other: &Self, - path: Option<&Path>, - writer: &mut dyn Write, - ) -> io::Result<()> { + /// Returns a diff between the original and modified source code. + /// + /// Returns `None` if `self` and `other` are not of the same kind. + pub fn diff<'a>( + &'a self, + other: &'a Self, + path: Option<&'a Path>, + ) -> Option> { match (self, other) { - (SourceKind::Python(src), SourceKind::Python(dst)) => { - let text_diff = TextDiff::from_lines(src, dst); - let mut unified_diff = text_diff.unified_diff(); + (SourceKind::Python(src), SourceKind::Python(dst)) => Some(SourceKindDiff { + kind: DiffKind::Python(src, dst), + path, + }), + (SourceKind::IpyNotebook(src), SourceKind::IpyNotebook(dst)) => Some(SourceKindDiff { + kind: DiffKind::IpyNotebook(src, dst), + path, + }), + _ => None, + } + } +} - if let Some(path) = path { - unified_diff.header(&fs::relativize_path(path), &fs::relativize_path(path)); - } +#[derive(Clone, Debug)] +pub struct SourceKindDiff<'a> { + kind: DiffKind<'a>, + path: Option<&'a Path>, +} - unified_diff.to_writer(&mut *writer)?; +impl std::fmt::Display for SourceKindDiff<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self.kind { + DiffKind::Python(original, modified) => { + let mut diff = CodeDiff::new(original, modified); - writer.write_all(b"\n")?; - writer.flush()?; + let relative_path = self.path.map(fs::relativize_path); - Ok(()) + if let Some(relative_path) = &relative_path { + diff.header(relative_path, relative_path); + } + + writeln!(f, "{diff}")?; } - (SourceKind::IpyNotebook(src), SourceKind::IpyNotebook(dst)) => { + DiffKind::IpyNotebook(original, modified) => { // Cell indices are 1-based. - for ((idx, src_cell), dst_cell) in - (1u32..).zip(src.cells().iter()).zip(dst.cells().iter()) + for ((idx, src_cell), dst_cell) in (1u32..) + .zip(original.cells().iter()) + .zip(modified.cells().iter()) { let (Cell::Code(src_cell), Cell::Code(dst_cell)) = (src_cell, dst_cell) else { continue; @@ -122,8 +145,18 @@ impl SourceKind { let src_source_code = src_cell.source.to_string(); let dst_source_code = dst_cell.source.to_string(); - let text_diff = TextDiff::from_lines(&src_source_code, &dst_source_code); - let mut unified_diff = text_diff.unified_diff(); + let header = self.path.map_or_else( + || (format!("cell {idx}"), format!("cell {idx}")), + |path| { + ( + format!("{}:cell {}", &fs::relativize_path(path), idx), + format!("{}:cell {}", &fs::relativize_path(path), idx), + ) + }, + ); + + let mut diff = CodeDiff::new(&src_source_code, &dst_source_code); + diff.header(&header.0, &header.1); // Jupyter notebook cells don't necessarily have a newline // at the end. For example, @@ -140,27 +173,85 @@ impl SourceKind { // print("hello") // // ``` - unified_diff.missing_newline_hint(false); + diff.missing_newline_hint(false); - if let Some(path) = path { - unified_diff.header( - &format!("{}:cell {}", &fs::relativize_path(path), idx), - &format!("{}:cell {}", &fs::relativize_path(path), idx), - ); - } else { - unified_diff.header(&format!("cell {idx}"), &format!("cell {idx}")); - }; - - unified_diff.to_writer(&mut *writer)?; + write!(f, "{diff}")?; } - writer.write_all(b"\n")?; - writer.flush()?; + writeln!(f)?; + } + } + + Ok(()) + } +} - Ok(()) +#[derive(Debug, Clone, Copy)] +enum DiffKind<'a> { + Python(&'a str, &'a str), + IpyNotebook(&'a Notebook, &'a Notebook), +} + +struct CodeDiff<'a> { + diff: TextDiff<'a, 'a, 'a, str>, + header: Option<(&'a str, &'a str)>, + missing_newline_hint: bool, +} + +impl<'a> CodeDiff<'a> { + fn new(original: &'a str, modified: &'a str) -> Self { + let diff = TextDiff::from_lines(original, modified); + Self { + diff, + header: None, + missing_newline_hint: true, + } + } + + fn header(&mut self, original: &'a str, modified: &'a str) { + self.header = Some((original, modified)); + } + + fn missing_newline_hint(&mut self, missing_newline_hint: bool) { + self.missing_newline_hint = missing_newline_hint; + } +} + +impl std::fmt::Display for CodeDiff<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some((original, modified)) = self.header { + writeln!(f, "--- {}", original.red())?; + writeln!(f, "+++ {}", modified.green())?; + } + + let mut unified = self.diff.unified_diff(); + unified.missing_newline_hint(self.missing_newline_hint); + + // Individual hunks (section of changes) + for hunk in unified.iter_hunks() { + writeln!(f, "{}", hunk.header())?; + + // individual lines + for change in hunk.iter_changes() { + match change.tag() { + ChangeTag::Equal => write!(f, " {}", change.value())?, + ChangeTag::Delete => write!(f, "{}{}", "-".red(), change.value())?, + ChangeTag::Insert => write!(f, "{}{}", "+".green(), change.value())?, + } + + if !self.diff.newline_terminated() { + writeln!(f)?; + } else if change.missing_newline() { + if self.missing_newline_hint { + writeln!(f, "{}", "\n\\ No newline at end of file".red())?; + } else { + writeln!(f)?; + } + } } - _ => panic!("cannot diff Python source code with Jupyter notebook source code"), } + + Ok(()) } }