[red-knot] Add 'Format document' to playground (#17217)

## Summary
This is more "because we can" than something we need. 

But since we're already building an "almost IDE" 

## Test Plan



https://github.com/user-attachments/assets/3a4bdad1-ba32-455a-9909-cfeb8caa1b28

Cargo.lock

@@ -2538,6 +2538,7 @@ dependencies = [
  "ruff_db",
  "ruff_macros",
  "ruff_python_ast",
+ "ruff_python_formatter",
  "ruff_text_size",
  "rustc-hash 2.1.1",
  "salsa",
@@ -2670,6 +2671,7 @@ dependencies = [
  "red_knot_python_semantic",
  "ruff_db",
  "ruff_notebook",
+ "ruff_python_formatter",
  "ruff_source_file",
  "ruff_text_size",
  "serde-wasm-bindgen",
@@ -3133,6 +3135,7 @@ dependencies = [
  "memchr",
  "regex",
  "ruff_cache",
+ "ruff_db",
  "ruff_formatter",
  "ruff_macros",
  "ruff_python_ast",
@@ -3141,6 +3144,7 @@ dependencies = [
  "ruff_source_file",
  "ruff_text_size",
  "rustc-hash 2.1.1",
+ "salsa",
  "schemars",
  "serde",
  "serde_json",

crates/red_knot_project/Cargo.toml

@@ -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/src/db.rs

@@ -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_wasm/Cargo.toml

@@ -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/src/lib.rs

@@ -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,10 @@ impl Workspace {
         Ok(format!("{:#?}", parsed.syntax()))
     }
 
+    pub fn format(&self, file_id: &FileHandle) -> Result<Option<String>, Error> {
+        formatted_file(&self.db, file_id.file).map_err(into_error)
+    }
+
     /// 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/ruff_db/src/files.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/parsed.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_python_formatter/Cargo.toml

@@ -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/src/db.rs

@@ -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/lib.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,28 @@ pub fn format_module_ast<'a>(
     Ok(formatted)
 }
 
+pub fn formatted_file(db: &dyn Db, file: File) -> Result<Option<String>, FormatModuleError> {
+    let options = db.format_options(file);
+
+    let parsed = parsed_module(db.upcast(), 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.upcast(), file);
+
+    let formatted = format_node(parsed, &comment_ranges, &source, options)?;
+    let printed = formatted.print()?;
+
+    if printed.as_code() == &*source {
+        Ok(None)
+    } else {
+        Ok(Some(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);

crates/ruff_python_parser/src/error.rs

@@ -7,7 +7,7 @@ use crate::TokenKind;
 
 /// Represents represent errors that occur during parsing and are
 /// returned by the `parse_*` functions.
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub struct ParseError {
     pub error: ParseErrorType,
     pub location: TextRange,
@@ -49,7 +49,7 @@ impl ParseError {
 }
 
 /// Represents the different types of errors that can occur during parsing of an f-string.
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum FStringErrorType {
     /// Expected a right brace after an opened left brace.
     UnclosedLbrace,
@@ -85,7 +85,7 @@ impl std::fmt::Display for FStringErrorType {
 }
 
 /// Represents the different types of errors that can occur during parsing.
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub enum ParseErrorType {
     /// An unexpected error occurred.
     OtherError(String),
@@ -362,7 +362,7 @@ impl std::fmt::Display for LexicalError {
 }
 
 /// Represents the different types of errors that can occur during lexing.
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum LexicalErrorType {
     // TODO: Can probably be removed, the places it is used seem to be able
     // to use the `UnicodeError` variant instead.

playground/knot/src/Editor/Editor.tsx

@@ -141,11 +141,13 @@ class PlaygroundServer
   implements
     languages.TypeDefinitionProvider,
     editor.ICodeEditorOpener,
-    languages.HoverProvider
+    languages.HoverProvider,
+    languages.DocumentFormattingEditProvider
 {
   private typeDefinitionProviderDisposable: IDisposable;
   private editorOpenerDisposable: IDisposable;
   private hoverDisposable: IDisposable;
+  private formatDisposable: IDisposable;
 
   constructor(
     private monaco: Monaco,
@@ -158,6 +160,8 @@ class PlaygroundServer
       this,
     );
     this.editorOpenerDisposable = monaco.editor.registerEditorOpener(this);
+    this.formatDisposable =
+      monaco.languages.registerDocumentFormattingEditProvider("python", this);
   }
 
   update(props: PlaygroundServerProps) {
@@ -353,10 +357,37 @@ class PlaygroundServer
     return true;
   }
 
+  provideDocumentFormattingEdits(
+    model: editor.ITextModel,
+  ): languages.ProviderResult<languages.TextEdit[]> {
+    if (this.props.files.selected == null) {
+      return null;
+    }
+
+    const fileHandle = this.props.files.handles[this.props.files.selected];
+
+    if (fileHandle == null) {
+      return null;
+    }
+
+    const formatted = this.props.workspace.format(fileHandle);
+    if (formatted != null) {
+      return [
+        {
+          range: model.getFullModelRange(),
+          text: formatted,
+        },
+      ];
+    }
+
+    return null;
+  }
+
   dispose() {
     this.hoverDisposable.dispose();
     this.editorOpenerDisposable.dispose();
     this.typeDefinitionProviderDisposable.dispose();
+    this.formatDisposable.dispose();
   }
 }