Make basic range formatting work in vs code

crates/ruff_cli/src/args.rs

@@ -11,6 +11,7 @@ use ruff_linter::settings::types::{
     FilePattern, PatternPrefixPair, PerFileIgnore, PreviewMode, PythonVersion, SerializationFormat,
 };
 use ruff_linter::{RuleParser, RuleSelector, RuleSelectorParser};
+use ruff_python_formatter::LspRowColumn;
 use ruff_workspace::configuration::{Configuration, RuleSelection};
 use ruff_workspace::resolver::ConfigurationTransformer;
 
@@ -395,6 +396,14 @@ pub struct FormatCommand {
     preview: bool,
     #[clap(long, overrides_with("preview"), hide = true)]
     no_preview: bool,
+    /// Range formatting start: Zero-indexed row and zero-indexed char-based column separated by
+    /// colon, e.g. `1:2`
+    #[clap(long)]
+    pub start: Option<LspRowColumn>,
+    /// Range formatting end: Zero-indexed row and zero-indexed char-based column separated by
+    /// colon, e.g. `3:4`
+    #[clap(long)]
+    pub end: Option<LspRowColumn>,
 }
 
 #[derive(Debug, Clone, Copy, clap::ValueEnum)]
@@ -516,6 +525,8 @@ impl FormatCommand {
                 files: self.files,
                 isolated: self.isolated,
                 stdin_filename: self.stdin_filename,
+                start: self.start,
+                end: self.end,
             },
             CliOverrides {
                 line_length: self.line_length,
@@ -572,6 +583,8 @@ pub struct FormatArguments {
     pub files: Vec<PathBuf>,
     pub isolated: bool,
     pub stdin_filename: Option<PathBuf>,
+    pub start: Option<LspRowColumn>,
+    pub end: Option<LspRowColumn>,
 }
 
 /// CLI settings that function as configuration overrides.

crates/ruff_cli/src/commands/format_stdin.rs

@@ -5,7 +5,9 @@ use anyhow::Result;
 use log::warn;
 
 use ruff_python_ast::PySourceType;
-use ruff_python_formatter::{format_module_source, PyFormatOptions};
+use ruff_python_formatter::{
+    format_module_source, format_module_source_range, LspRowColumn, PyFormatOptions,
+};
 use ruff_workspace::resolver::python_file_at_path;
 
 use crate::args::{CliOverrides, FormatArguments};
@@ -42,7 +44,7 @@ pub(crate) fn format_stdin(cli: &FormatArguments, overrides: &CliOverrides) -> R
         .formatter
         .to_format_options(path.map(PySourceType::from).unwrap_or_default());
 
-    match format_source(path, options, mode) {
+    match format_source(path, options, mode, cli.start, cli.end) {
         Ok(result) => match mode {
             FormatMode::Write => Ok(ExitStatus::Success),
             FormatMode::Check => {
@@ -65,12 +67,21 @@ fn format_source(
     path: Option<&Path>,
     options: PyFormatOptions,
     mode: FormatMode,
+    start: Option<LspRowColumn>,
+    end: Option<LspRowColumn>,
 ) -> Result<FormatCommandResult, FormatCommandError> {
     let unformatted = read_from_stdin()
         .map_err(|err| FormatCommandError::Read(path.map(Path::to_path_buf), err))?;
+    let formatted = if start.is_some() || end.is_some() {
+        let formatted = format_module_source_range(&unformatted, options, start, end)
+            .map_err(|err| FormatCommandError::FormatModule(path.map(Path::to_path_buf), err))?;
+        formatted
+    } else {
         let formatted = format_module_source(&unformatted, options)
             .map_err(|err| FormatCommandError::FormatModule(path.map(Path::to_path_buf), err))?;
         let formatted = formatted.as_code();
+        formatted.to_string()
+    };
     if formatted.len() == unformatted.len() && formatted == unformatted {
         Ok(FormatCommandResult::Unchanged)
     } else {

crates/ruff_formatter/src/diagnostics.rs

@@ -1,6 +1,5 @@
 use crate::prelude::TagKind;
 use crate::GroupId;
-use ruff_text_size::TextRange;
 use std::error::Error;
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
@@ -12,7 +11,7 @@ pub enum FormatError {
     SyntaxError { message: &'static str },
     /// In case range formatting failed because the provided range was larger
     /// than the formatted syntax tree
-    RangeError { input: TextRange, tree: TextRange },
+    RangeError { row: usize, col: usize },
 
     /// In case printing the document failed because it has an invalid structure.
     InvalidDocument(InvalidDocumentError),
@@ -32,9 +31,9 @@ impl std::fmt::Display for FormatError {
             FormatError::SyntaxError {message} => {
                 std::write!(fmt, "syntax error: {message}")
             },
-            FormatError::RangeError { input, tree } => std::write!(
+            FormatError::RangeError { row, col } => std::write!(
                 fmt,
-                "formatting range {input:?} is larger than syntax tree {tree:?}"
+                "formatting range {row}:{col} is not a valid index"
             ),
             FormatError::InvalidDocument(error) => std::write!(fmt, "Invalid document: {error}\n\n This is an internal Rome error. Please report if necessary."),
             FormatError::PoorLayout => {

crates/ruff_python_formatter/src/cli.rs

@@ -8,10 +8,11 @@ use clap::{command, Parser, ValueEnum};
 use ruff_formatter::SourceCode;
 use ruff_python_index::tokens_and_ranges;
 use ruff_python_parser::{parse_ok_tokens, Mode};
-use ruff_text_size::Ranged;
+use ruff_source_file::Locator;
+use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
 use crate::comments::collect_comments;
-use crate::{format_module_ast, PyFormatOptions};
+use crate::{format_module_ast, format_module_range, PyFormatOptions};
 
 #[derive(ValueEnum, Clone, Debug)]
 pub enum Emit {
@@ -36,19 +37,38 @@ pub struct Cli {
     pub print_ir: bool,
     #[clap(long)]
     pub print_comments: bool,
+    /// byte offset for range formatting
+    #[clap(long)]
+    pub start: Option<u32>,
+    /// byte offset for range formatting
+    #[clap(long)]
+    pub end: Option<u32>,
 }
 
 pub fn format_and_debug_print(source: &str, cli: &Cli, source_type: &Path) -> Result<String> {
     let (tokens, comment_ranges) = tokens_and_ranges(source)
         .map_err(|err| format_err!("Source contains syntax errors {err:?}"))?;
-
-    // Parse the AST.
     let module =
         parse_ok_tokens(tokens, Mode::Module, "<filename>").context("Syntax error in input")?;
-
     let options = PyFormatOptions::from_extension(source_type);
-
     let source_code = SourceCode::new(source);
+    let locator = Locator::new(source);
+
+    if cli.start.is_some() || cli.end.is_some() {
+        let range = TextRange::new(
+            cli.start.map(TextSize::new).unwrap_or_default(),
+            cli.end.map(TextSize::new).unwrap_or(source.text_len()),
+        );
+        return Ok(format_module_range(
+            &module,
+            &comment_ranges,
+            source,
+            options,
+            &locator,
+            range,
+        )?);
+    }
+
     let formatted = format_module_ast(&module, &comment_ranges, source, options)
         .context("Failed to format node")?;
     if cli.print_ir {

crates/ruff_python_formatter/src/lib.rs

@@ -1,5 +1,6 @@
+use std::str::FromStr;
 use thiserror::Error;
-use tracing::Level;
+use tracing::{warn, Level};
 
 use ruff_formatter::prelude::*;
 use ruff_formatter::{format, FormatError, Formatted, PrintError, Printed, SourceCode};
@@ -10,12 +11,14 @@ use ruff_python_parser::lexer::LexicalError;
 use ruff_python_parser::{parse_ok_tokens, Mode, ParseError};
 use ruff_python_trivia::CommentRanges;
 use ruff_source_file::Locator;
+use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
 use crate::comments::{
     dangling_comments, leading_comments, trailing_comments, Comments, SourceComment,
 };
 pub use crate::context::PyFormatContext;
 pub use crate::options::{MagicTrailingComma, PreviewMode, PyFormatOptions, QuoteStyle};
+use crate::statement::suite::SuiteKind;
 use crate::verbatim::suppressed_node;
 pub use settings::FormatterSettings;
 
@@ -134,6 +137,73 @@ pub fn format_module_source(
     Ok(formatted.print()?)
 }
 
+/// Range formatting coordinate: Zero-indexed row and zero-indexed char-based column separated by
+/// colon, e.g. `1:2`.
+///
+/// See [`Locator::convert_row_and_column`] for details on the semantics.
+#[derive(Copy, Clone, Debug, Default)]
+pub struct LspRowColumn {
+    row: usize,
+    col: usize,
+}
+
+impl FromStr for LspRowColumn {
+    type Err = &'static str;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let Some((row, col)) = s.split_once(':') else {
+            return Err("Coordinate is missing a colon, the format is `<row>:<column>`");
+        };
+
+        Ok(LspRowColumn {
+            row: row.parse().map_err(|_| "row must be a number")?,
+            col: col.parse().map_err(|_| "col must be a number")?,
+        })
+    }
+}
+#[tracing::instrument(name = "format", level = Level::TRACE, skip_all)]
+pub fn format_module_source_range(
+    source: &str,
+    options: PyFormatOptions,
+    start: Option<LspRowColumn>,
+    end: Option<LspRowColumn>,
+) -> Result<String, FormatModuleError> {
+    let (tokens, comment_ranges) = tokens_and_ranges(source)?;
+    let module = parse_ok_tokens(tokens, Mode::Module, "<filename>")?;
+    let locator = Locator::new(source);
+
+    let start = if let Some(start) = start {
+        locator
+            .convert_row_and_column(start.row, start.col)
+            .ok_or(FormatError::RangeError {
+                row: start.row,
+                col: start.col,
+            })?
+    } else {
+        TextSize::default()
+    };
+    let end = if let Some(end) = end {
+        locator
+            .convert_row_and_column(end.row, end.col)
+            .ok_or(FormatError::RangeError {
+                row: end.row,
+                col: end.col,
+            })?
+    } else {
+        source.text_len()
+    };
+
+    let formatted = format_module_range(
+        &module,
+        &comment_ranges,
+        source,
+        options,
+        &locator,
+        TextRange::new(start, end),
+    )?;
+    Ok(formatted)
+}
+
 pub fn format_module_ast<'a>(
     module: &'a Mod,
     comment_ranges: &'a CommentRanges,
@@ -155,6 +225,57 @@ pub fn format_module_ast<'a>(
     Ok(formatted)
 }
 
+pub fn format_module_range<'a>(
+    module: &'a Mod,
+    comment_ranges: &'a CommentRanges,
+    source: &'a str,
+    options: PyFormatOptions,
+    locator: &Locator<'a>,
+    range: TextRange,
+) -> FormatResult<String> {
+    let comments = Comments::from_ast(&module, SourceCode::new(source), &comment_ranges);
+
+    let Mod::Module(module_inner) = &module else {
+        panic!("That's not a module");
+    };
+
+    // ```
+    // a = 1; b = 2; c = 3; d = 4; e = 5
+    //             ^ b end  ^ d start
+    //          ^^^^^^^^^^^^^^^ range
+    //          ^ range start ^ range end
+    // ```
+    // TODO: If it goes beyond the end of the last stmt or before start, do we need to format the
+    // parent?
+    // TODO: Change suite formatting so we can use a slice instead
+    let in_range: Vec<_> = module_inner
+        .body
+        .iter()
+        .cloned()
+        // TODO: check whether these bounds need equality
+        .skip_while(|child| range.start() > child.end())
+        .take_while(|child| child.start() < range.end())
+        .collect();
+    let (Some(first), Some(last)) = (in_range.first(), in_range.last()) else {
+        // TODO: Use tracing again https://github.com/tokio-rs/tracing/issues/2721
+        // TODO: Forward this to something proper
+        eprintln!("The formatting range contains no statements");
+        return Ok(source.to_string());
+    };
+
+    let mut buffer = source[TextRange::up_to(first.start())].to_string();
+
+    let formatted: Formatted<PyFormatContext> = format!(
+        PyFormatContext::new(options.clone(), locator.contents(), comments),
+        [in_range.format().with_options(SuiteKind::TopLevel)]
+    )?;
+    //println!("{}", formatted.document().display(SourceCode::new(source)));
+    // TODO: Make the printer use the buffer instead
+    buffer += formatted.print()?.as_code();
+    buffer += &source[TextRange::new(last.end(), source.text_len())];
+    return Ok(buffer.to_string());
+}
+
 /// Public function for generating a printable string of the debug comments.
 pub fn pretty_comments(module: &Mod, comment_ranges: &CommentRanges, source: &str) -> String {
     let source_code = SourceCode::new(source);