Add formatter to runes

Cargo.lock

@@ -1835,9 +1835,11 @@ dependencies = [
  "notify",
  "parking_lot",
  "rayon",
+ "ruff_formatter",
  "ruff_index",
  "ruff_notebook",
  "ruff_python_ast",
+ "ruff_python_formatter",
  "ruff_python_parser",
  "ruff_python_trivia",
  "ruff_text_size",

crates/red_knot/Cargo.toml

@@ -12,12 +12,14 @@ license.workspace = true
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-ruff_python_parser = { path = "../ruff_python_parser" }
-ruff_python_ast = { path = "../ruff_python_ast" }
-ruff_python_trivia = { path = "../ruff_python_trivia" }
-ruff_text_size = { path = "../ruff_text_size" }
+ruff_formatter = { path = "../ruff_formatter" }
 ruff_index = { path = "../ruff_index" }
 ruff_notebook = { path = "../ruff_notebook" }
+ruff_python_ast = { path = "../ruff_python_ast" }
+ruff_python_formatter = { path = "../ruff_python_formatter" }
+ruff_python_parser = { path = "../ruff_python_parser" }
+ruff_python_trivia = { path = "../ruff_python_trivia" }
+ruff_text_size = { path = "../ruff_text_size" }
 
 anyhow = { workspace = true }
 bitflags = { workspace = true }

crates/red_knot/src/format.rs

@@ -0,0 +1,135 @@
+use std::ops::{Deref, DerefMut};
+
+use ruff_formatter::PrintedRange;
+use ruff_python_formatter::{FormatModuleError, PyFormatOptions};
+use ruff_text_size::TextRange;
+
+use crate::cache::KeyValueCache;
+use crate::db::{HasJar, QueryError, SourceDb};
+use crate::files::FileId;
+use crate::lint::Diagnostics;
+use crate::FxDashSet;
+
+pub(crate) trait FormatDb: SourceDb {
+    /// Formats a file and returns its formatted content or an indicator that it is unchanged.
+    fn format_file(&self, file_id: FileId) -> Result<FormattedFile, FormatError>;
+
+    /// Formats a range in a file.
+    fn format_file_range(
+        &self,
+        file_id: FileId,
+        range: TextRange,
+    ) -> Result<PrintedRange, FormatError>;
+
+    fn check_file_formatted(&self, file_id: FileId) -> Result<Diagnostics, FormatError>;
+}
+
+#[tracing::instrument(level = "trace", skip(db))]
+pub(crate) fn format_file<Db>(db: &Db, file_id: FileId) -> Result<FormattedFile, FormatError>
+where
+    Db: FormatDb + HasJar<FormatJar>,
+{
+    let formatted = &db.jar()?.formatted;
+
+    if formatted.contains(&file_id) {
+        return Ok(FormattedFile::Unchanged);
+    }
+
+    let source = db.source(file_id)?;
+
+    // TODO use the `format_module` method here to re-use the AST.
+    let printed =
+        ruff_python_formatter::format_module_source(source.text(), PyFormatOptions::default())?;
+
+    Ok(if printed.as_code() == source.text() {
+        formatted.insert(file_id);
+        FormattedFile::Unchanged
+    } else {
+        FormattedFile::Formatted(printed.into_code())
+    })
+}
+
+#[tracing::instrument(level = "trace", skip(db))]
+pub(crate) fn format_file_range<Db: FormatDb + HasJar<FormatJar>>(
+    db: &Db,
+    file_id: FileId,
+    range: TextRange,
+) -> Result<PrintedRange, FormatError> {
+    let formatted = &db.jar()?.formatted;
+    let source = db.source(file_id)?;
+
+    if formatted.contains(&file_id) {
+        return Ok(PrintedRange::new(source.text()[range].into(), range));
+    }
+
+    // TODO use the `format_module` method here to re-use the AST.
+
+    let result =
+        ruff_python_formatter::format_range(source.text(), range, PyFormatOptions::default())?;
+    Ok(result)
+}
+
+/// Checks if the file is correctly formatted. It creates a diagnostic for formatting issues.
+#[tracing::instrument(level = "trace", skip(db))]
+pub(crate) fn check_formatted<Db>(db: &Db, file_id: FileId) -> Result<Diagnostics, FormatError>
+where
+    Db: FormatDb + HasJar<FormatJar>,
+{
+    Ok(if db.format_file(file_id)?.is_unchanged() {
+        Diagnostics::Empty
+    } else {
+        Diagnostics::from(vec!["File is not formatted".to_string()])
+    })
+}
+
+#[derive(Debug)]
+pub(crate) enum FormatError {
+    Format(FormatModuleError),
+    Query(QueryError),
+}
+
+impl From<FormatModuleError> for FormatError {
+    fn from(value: FormatModuleError) -> Self {
+        Self::Format(value)
+    }
+}
+
+impl From<QueryError> for FormatError {
+    fn from(value: QueryError) -> Self {
+        Self::Query(value)
+    }
+}
+
+#[derive(Clone, Eq, PartialEq, Debug)]
+pub(crate) enum FormattedFile {
+    Formatted(String),
+    Unchanged,
+}
+
+impl FormattedFile {
+    pub(crate) const fn is_unchanged(&self) -> bool {
+        matches!(self, FormattedFile::Unchanged)
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct FormatJar {
+    pub formatted: FxDashSet<FileId>,
+}
+
+#[derive(Default, Debug)]
+pub(crate) struct FormattedStorage(KeyValueCache<FileId, ()>);
+
+impl Deref for FormattedStorage {
+    type Target = KeyValueCache<FileId, ()>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for FormattedStorage {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}

crates/red_knot/src/lib.rs

@@ -12,6 +12,7 @@ pub mod cache;
 pub mod cancellation;
 pub mod db;
 pub mod files;
+mod format;
 pub mod hir;
 pub mod lint;
 pub mod module;

crates/red_knot/src/main.rs

@@ -13,7 +13,9 @@ use tracing_subscriber::layer::{Context, Filter, SubscriberExt};
 use tracing_subscriber::{Layer, Registry};
 use tracing_tree::time::Uptime;
 
-use red_knot::db::{HasJar, ParallelDatabase, QueryError, SemanticDb, SourceDb, SourceJar};
+use red_knot::db::{
+    Database, HasJar, ParallelDatabase, QueryError, SemanticDb, SourceDb, SourceJar,
+};
 use red_knot::files::FileId;
 use red_knot::module::{ModuleSearchPath, ModuleSearchPathKind};
 use red_knot::program::check::ExecutionMode;
@@ -138,22 +140,28 @@ impl MainLoop {
 
             match message {
                 MainLoopMessage::CheckProgram { revision } => {
-                    let program = program.snapshot();
-                    let sender = self.orchestrator_sender.clone();
+                    {
+                        let program = program.snapshot();
+                        let sender = self.orchestrator_sender.clone();
 
-                    // Spawn a new task that checks the program. This needs to be done in a separate thread
-                    // to prevent blocking the main loop here.
-                    rayon::spawn(move || match program.check(ExecutionMode::ThreadPool) {
-                        Ok(result) => {
-                            sender
-                                .send(OrchestratorMessage::CheckProgramCompleted {
-                                    diagnostics: result,
-                                    revision,
-                                })
-                                .unwrap();
-                        }
-                        Err(QueryError::Cancelled) => {}
-                    });
+                        // Spawn a new task that checks the program. This needs to be done in a separate thread
+                        // to prevent blocking the main loop here.
+                        rayon::spawn(move || match program.check(ExecutionMode::ThreadPool) {
+                            Ok(result) => {
+                                sender
+                                    .send(OrchestratorMessage::CheckProgramCompleted {
+                                        diagnostics: result,
+                                        revision,
+                                    })
+                                    .unwrap();
+                            }
+                            Err(QueryError::Cancelled) => {}
+                        });
+                    }
+
+                    if !program.is_cancelled() {
+                        let _ = program.format();
+                    }
                 }
                 MainLoopMessage::ApplyChanges(changes) => {
                     // Automatically cancels any pending queries and waits for them to complete.

crates/red_knot/src/program/check.rs

@@ -3,6 +3,7 @@ use rustc_hash::FxHashSet;
 
 use crate::db::{Database, LintDb, QueryError, QueryResult, SemanticDb};
 use crate::files::FileId;
+use crate::format::{FormatDb, FormatError};
 use crate::lint::Diagnostics;
 use crate::program::Program;
 use crate::symbols::Dependency;
@@ -64,6 +65,18 @@ impl Program {
         if self.workspace().is_file_open(file) {
             diagnostics.extend_from_slice(&self.lint_syntax(file)?);
             diagnostics.extend_from_slice(&self.lint_semantic(file)?);
+
+            match self.check_file_formatted(file) {
+                Ok(format_diagnostics) => {
+                    diagnostics.extend_from_slice(&format_diagnostics);
+                }
+                Err(FormatError::Query(err)) => {
+                    return Err(err);
+                }
+                Err(FormatError::Format(error)) => {
+                    diagnostics.push(format!("Error formatting file: {error}"));
+                }
+            }
         }
 
         Ok(Diagnostics::from(diagnostics))

crates/red_knot/src/program/format.rs

@@ -0,0 +1,44 @@
+use crate::db::{QueryResult, SourceDb};
+use crate::format::{FormatDb, FormatError, FormattedFile};
+use crate::program::Program;
+
+impl Program {
+    #[tracing::instrument(level = "trace", skip(self))]
+    pub fn format(&mut self) -> QueryResult<()> {
+        // Formats all open files
+
+        // TODO make `Executor` from `check` reusable.
+        for file in self.workspace.open_files() {
+            match self.format_file(file) {
+                Ok(FormattedFile::Formatted(content)) => {
+                    let path = self.file_path(file);
+
+                    // TODO: This is problematic because it immediately re-triggers the file watcher.
+                    //  A possible solution is to track the self "inflicted" changes inside of programs
+                    //  by tracking the file revision right after the write. It could then use the revision
+                    //  to determine which changes are safe to ignore (and in which context).
+                    //  An other alternative is to not write as part of the `format` command and instead
+                    //  return a Vec with the format results and leave the writing to the caller.
+                    //  I think that's undesired because a) we still need a way to tell the formatter
+                    //  that it won't be necessary to format the content again and
+                    //  b) it would reduce concurrency because the writing would need to wait for the file
+                    //  formatting to be complete, unless we use some form of communication channel.
+                    std::fs::write(path, content).expect("Unable to write file");
+                }
+                Ok(FormattedFile::Unchanged) => {
+                    // No op
+                }
+                Err(FormatError::Query(error)) => {
+                    return Err(error);
+                }
+                Err(FormatError::Format(error)) => {
+                    // TODO proper error handling. We should either propagate this error or
+                    //  emit a diagnostic (probably this).
+                    tracing::warn!("Failed to format file: {}", error);
+                }
+            }
+        }
+
+        Ok(())
+    }
+}

crates/red_knot/src/program/mod.rs

@@ -1,3 +1,5 @@
+use ruff_formatter::PrintedRange;
+use ruff_text_size::TextRange;
 use std::path::Path;
 use std::sync::Arc;
 
@@ -6,6 +8,10 @@ use crate::db::{
     QueryResult, SemanticDb, SemanticJar, Snapshot, SourceDb, SourceJar,
 };
 use crate::files::{FileId, Files};
+use crate::format::{
+    check_formatted, format_file, format_file_range, FormatDb, FormatError, FormatJar,
+    FormattedFile,
+};
 use crate::lint::{lint_semantic, lint_syntax, Diagnostics};
 use crate::module::{
     add_module, file_to_module, path_to_module, resolve_module, set_module_search_paths, Module,
@@ -18,6 +24,7 @@ use crate::types::{infer_symbol_type, Type};
 use crate::Workspace;
 
 pub mod check;
+mod format;
 
 #[derive(Debug)]
 pub struct Program {
@@ -39,7 +46,7 @@ impl Program {
     where
         I: IntoIterator<Item = FileChange>,
     {
-        let (source, semantic, lint) = self.jars_mut();
+        let (source, semantic, lint, format) = self.jars_mut();
         for change in changes {
             semantic.module_resolver.remove_module(change.id);
             semantic.symbol_tables.remove(&change.id);
@@ -49,6 +56,7 @@ impl Program {
             semantic.type_store.remove_module(change.id);
             lint.lint_syntax.remove(&change.id);
             lint.lint_semantic.remove(&change.id);
+            format.formatted.remove(&change.id);
         }
     }
 
@@ -124,6 +132,24 @@ impl LintDb for Program {
     }
 }
 
+impl FormatDb for Program {
+    fn format_file(&self, file_id: FileId) -> Result<FormattedFile, FormatError> {
+        format_file(self, file_id)
+    }
+
+    fn format_file_range(
+        &self,
+        file_id: FileId,
+        range: TextRange,
+    ) -> Result<PrintedRange, FormatError> {
+        format_file_range(self, file_id, range)
+    }
+
+    fn check_file_formatted(&self, file_id: FileId) -> Result<Diagnostics, FormatError> {
+        check_formatted(self, file_id)
+    }
+}
+
 impl Db for Program {}
 
 impl Database for Program {
@@ -147,7 +173,7 @@ impl ParallelDatabase for Program {
 }
 
 impl HasJars for Program {
-    type Jars = (SourceJar, SemanticJar, LintJar);
+    type Jars = (SourceJar, SemanticJar, LintJar, FormatJar);
 
     fn jars(&self) -> QueryResult<&Self::Jars> {
         self.jars.jars()
@@ -188,6 +214,16 @@ impl HasJar<LintJar> for Program {
     }
 }
 
+impl HasJar<FormatJar> for Program {
+    fn jar(&self) -> QueryResult<&FormatJar> {
+        Ok(&self.jars()?.3)
+    }
+
+    fn jar_mut(&mut self) -> &mut FormatJar {
+        &mut self.jars_mut().3
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 pub struct FileChange {
     id: FileId,