[red-knot] Add 'Format document' to playground

Author: MichaReiser

Cargo.lock

@@ -16,6 +16,7 @@ ruff_cache = { workspace = true }
 ruff_db = { workspace = true, features = ["cache", "serde"] }
 ruff_macros = { workspace = true }
 ruff_python_ast = { workspace = true, features = ["serde"] }
+ruff_python_formatter = { workspace = true, optional = true }
 ruff_text_size = { workspace = true }
 red_knot_ide = { workspace = true }
 red_knot_python_semantic = { workspace = true, features = ["serde"] }
@@ -43,8 +44,13 @@ insta = { workspace = true, features = ["redactions", "ron"] }
 [features]
 default = ["zstd"]
 deflate = ["red_knot_vendored/deflate"]
-schemars = ["dep:schemars", "ruff_db/schemars", "red_knot_python_semantic/schemars"]
+schemars = [
+    "dep:schemars",
+    "ruff_db/schemars",
+    "red_knot_python_semantic/schemars",
+]
 zstd = ["red_knot_vendored/zstd"]
+format = ["ruff_python_formatter"]
 
 [lints]
 workspace = true

crates/red_knot_project/Cargo.toml

@@ -174,6 +174,32 @@ impl Db for ProjectDatabase {
     }
 }
 
+#[cfg(feature = "format")]
+mod format {
+    use crate::ProjectDatabase;
+    use ruff_db::files::File;
+    use ruff_db::Upcast;
+    use ruff_python_formatter::{Db as FormatDb, PyFormatOptions};
+
+    #[salsa::db]
+    impl FormatDb for ProjectDatabase {
+        fn format_options(&self, file: File) -> PyFormatOptions {
+            let source_ty = file.source_type(self);
+            PyFormatOptions::from_source_type(source_ty)
+        }
+    }
+
+    impl Upcast<dyn FormatDb> for ProjectDatabase {
+        fn upcast(&self) -> &(dyn FormatDb + 'static) {
+            self
+        }
+
+        fn upcast_mut(&mut self) -> &mut (dyn FormatDb + 'static) {
+            self
+        }
+    }
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
     use std::sync::Arc;

crates/red_knot_project/src/db.rs

@@ -19,12 +19,16 @@ doctest = false
 default = ["console_error_panic_hook"]
 
 [dependencies]
-red_knot_project = { workspace = true, default-features = false, features = ["deflate"] }
 red_knot_ide = { workspace = true }
+red_knot_project = { workspace = true, default-features = false, features = [
+    "deflate",
+    "format"
+] }
 red_knot_python_semantic = { workspace = true }
 
 ruff_db = { workspace = true, default-features = false, features = [] }
 ruff_notebook = { workspace = true }
+ruff_python_formatter = { workspace = true }
 ruff_source_file = { workspace = true }
 ruff_text_size = { workspace = true }
 
@@ -44,4 +48,3 @@ wasm-bindgen-test = { workspace = true }
 
 [lints]
 workspace = true
-

crates/red_knot_wasm/Cargo.toml

@@ -18,6 +18,7 @@ use ruff_db::system::{
 };
 use ruff_db::Upcast;
 use ruff_notebook::Notebook;
+use ruff_python_formatter::formatted_file;
 use ruff_source_file::{LineIndex, OneIndexed, SourceLocation};
 use ruff_text_size::Ranged;
 use wasm_bindgen::prelude::*;
@@ -142,7 +143,11 @@ impl Workspace {
     }
 
     #[wasm_bindgen(js_name = "closeFile")]
-    pub fn close_file(&mut self, file_id: &FileHandle) -> Result<(), Error> {
+    #[allow(
+        clippy::needless_pass_by_value,
+        reason = "It's intentional that the file handle is consumed because it is no longer valid after closing"
+    )]
+    pub fn close_file(&mut self, file_id: FileHandle) -> Result<(), Error> {
         let file = file_id.file;
 
         self.db.project().close_file(&mut self.db, file);
@@ -184,6 +189,14 @@ impl Workspace {
         Ok(format!("{:#?}", parsed.syntax()))
     }
 
+    pub fn format(&self, file_id: &FileHandle) -> Result<Option<String>, Error> {
+        let formatted = formatted_file(&self.db, file_id.file)
+            .as_ref()
+            .map_err(into_error)?;
+
+        Ok(formatted.clone().into_reformatted())
+    }
+
     /// Returns the token stream for `path` serialized as a string.
     pub fn tokens(&self, file_id: &FileHandle) -> Result<String, Error> {
         let parsed = ruff_db::parsed::parsed_module(&self.db, file_id.file);

crates/red_knot_wasm/src/lib.rs

@@ -424,9 +424,19 @@ impl File {
 
     /// Returns `true` if the file should be analyzed as a type stub.
     pub fn is_stub(self, db: &dyn Db) -> bool {
-        self.path(db)
-            .extension()
-            .is_some_and(|extension| PySourceType::from_extension(extension).is_stub())
+        self.source_type(db).is_stub()
+    }
+
+    pub fn source_type(self, db: &dyn Db) -> PySourceType {
+        match self.path(db) {
+            FilePath::System(path) => path
+                .extension()
+                .map_or(PySourceType::Python, PySourceType::from_extension),
+            FilePath::Vendored(_) => PySourceType::Stub,
+            FilePath::SystemVirtual(path) => path
+                .extension()
+                .map_or(PySourceType::Python, PySourceType::from_extension),
+        }
     }
 }
 

crates/ruff_db/src/files.rs

@@ -2,10 +2,10 @@ use std::fmt::Formatter;
 use std::ops::Deref;
 use std::sync::Arc;
 
-use ruff_python_ast::{ModModule, PySourceType};
+use ruff_python_ast::ModModule;
 use ruff_python_parser::{parse_unchecked_source, Parsed};
 
-use crate::files::{File, FilePath};
+use crate::files::File;
 use crate::source::source_text;
 use crate::Db;
 
@@ -25,17 +25,7 @@ pub fn parsed_module(db: &dyn Db, file: File) -> ParsedModule {
     let _span = tracing::trace_span!("parsed_module", ?file).entered();
 
     let source = source_text(db, file);
-    let path = file.path(db);
-
-    let ty = match path {
-        FilePath::System(path) => path
-            .extension()
-            .map_or(PySourceType::Python, PySourceType::from_extension),
-        FilePath::Vendored(_) => PySourceType::Stub,
-        FilePath::SystemVirtual(path) => path
-            .extension()
-            .map_or(PySourceType::Python, PySourceType::from_extension),
-    };
+    let ty = file.source_type(db);
 
     ParsedModule::new(parse_unchecked_source(&source, ty))
 }

crates/ruff_db/src/parsed.rs

@@ -15,6 +15,7 @@ doctest = false
 
 [dependencies]
 ruff_cache = { workspace = true }
+ruff_db = { workspace = true }
 ruff_formatter = { workspace = true }
 ruff_macros = { workspace = true }
 ruff_python_trivia = { workspace = true }
@@ -30,6 +31,7 @@ itertools = { workspace = true }
 memchr = { workspace = true }
 regex = { workspace = true }
 rustc-hash = { workspace = true }
+salsa = { workspace = true }
 serde = { workspace = true, optional = true }
 schemars = { workspace = true, optional = true }
 smallvec = { workspace = true }
@@ -58,7 +60,12 @@ required-features = ["serde"]
 
 [features]
 default = ["serde"]
-serde = ["dep:serde", "ruff_formatter/serde", "ruff_source_file/serde", "ruff_python_ast/serde"]
+serde = [
+    "dep:serde",
+    "ruff_formatter/serde",
+    "ruff_source_file/serde",
+    "ruff_python_ast/serde",
+]
 schemars = ["dep:schemars", "ruff_formatter/schemars"]
 
 [lints]

crates/ruff_python_formatter/Cargo.toml

@@ -0,0 +1,9 @@
+use ruff_db::{files::File, Db as SourceDb, Upcast};
+
+use crate::PyFormatOptions;
+
+#[salsa::db]
+pub trait Db: SourceDb + Upcast<dyn SourceDb> {
+    /// Returns the formatting options
+    fn format_options(&self, file: File) -> PyFormatOptions;
+}

crates/ruff_python_formatter/src/db.rs

@@ -1,3 +1,6 @@
+use ruff_db::files::File;
+use ruff_db::parsed::parsed_module;
+use ruff_db::source::source_text;
 use thiserror::Error;
 use tracing::Level;
 
@@ -13,6 +16,7 @@ use crate::comments::{
     has_skip_comment, leading_comments, trailing_comments, Comments, SourceComment,
 };
 pub use crate::context::PyFormatContext;
+pub use crate::db::Db;
 pub use crate::options::{
     DocstringCode, DocstringCodeLineWidth, MagicTrailingComma, PreviewMode, PyFormatOptions,
     QuoteStyle,
@@ -25,6 +29,7 @@ pub(crate) mod builders;
 pub mod cli;
 mod comments;
 pub(crate) mod context;
+mod db;
 pub(crate) mod expression;
 mod generated;
 pub(crate) mod module;
@@ -96,7 +101,7 @@ where
     }
 }
 
-#[derive(Error, Debug)]
+#[derive(Error, Debug, salsa::Update, PartialEq, Eq)]
 pub enum FormatModuleError {
     #[error(transparent)]
     ParseError(#[from] ParseError),
@@ -124,6 +129,19 @@ pub fn format_module_ast<'a>(
     source: &'a str,
     options: PyFormatOptions,
 ) -> FormatResult<Formatted<PyFormatContext<'a>>> {
+    format_node(parsed, comment_ranges, source, options)
+}
+
+fn format_node<'a, N>(
+    parsed: &'a Parsed<N>,
+    comment_ranges: &'a CommentRanges,
+    source: &'a str,
+    options: PyFormatOptions,
+) -> FormatResult<Formatted<PyFormatContext<'a>>>
+where
+    N: AsFormat<PyFormatContext<'a>>,
+    &'a N: Into<AnyNodeRef<'a>>,
+{
     let source_code = SourceCode::new(source);
     let comments = Comments::from_ast(parsed.syntax(), source_code, comment_ranges);
 
@@ -138,6 +156,29 @@ pub fn format_module_ast<'a>(
     Ok(formatted)
 }
 
+#[salsa::tracked(return_ref)]
+pub fn formatted_file(db: &dyn Db, file: File) -> Result<FormattedFile, FormatModuleError> {
+    let options = db.format_options(file);
+
+    let parsed = parsed_module(db, file);
+
+    if let Some(first) = parsed.errors().first() {
+        return Err(FormatModuleError::ParseError(first.clone()));
+    }
+
+    let comment_ranges = CommentRanges::from(parsed.tokens());
+    let source = source_text(db, file);
+
+    let formatted = format_node(parsed, &comment_ranges, &source, options)?;
+    let printed = formatted.print()?;
+
+    if printed.as_code() == &*source {
+        Ok(FormattedFile::Formatted)
+    } else {
+        Ok(FormattedFile::Reformatted(printed.into_code()))
+    }
+}
+
 /// 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);
@@ -146,6 +187,31 @@ pub fn pretty_comments(module: &Mod, comment_ranges: &CommentRanges, source: &st
     std::format!("{comments:#?}", comments = comments.debug(source_code))
 }
 
+#[derive(Debug, Clone, Eq, PartialEq, salsa::Update)]
+pub enum FormattedFile {
+    /// The file is already correctly formatted.
+    Formatted,
+
+    /// The file has been reformatted. The string is its newly formatted content.
+    Reformatted(String),
+}
+
+impl FormattedFile {
+    pub fn as_reformatted(&self) -> Option<&str> {
+        match self {
+            FormattedFile::Formatted => None,
+            FormattedFile::Reformatted(code) => Some(code),
+        }
+    }
+
+    pub fn into_reformatted(self) -> Option<String> {
+        match self {
+            FormattedFile::Formatted => None,
+            FormattedFile::Reformatted(code) => Some(code),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::path::Path;