[ty] Implement streaming for workspace diagnostics

Author: MichaReiser

crates/ty/src/lib.rs

@@ -22,12 +22,13 @@ use colored::Colorize;
 use crossbeam::channel as crossbeam_channel;
 use rayon::ThreadPoolBuilder;
 use ruff_db::diagnostic::{Diagnostic, DisplayDiagnosticConfig, Severity};
+use ruff_db::files::File;
 use ruff_db::max_parallelism;
 use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
 use salsa::plumbing::ZalsaDatabase;
 use ty_project::metadata::options::ProjectOptionsOverrides;
 use ty_project::watch::ProjectWatcher;
-use ty_project::{Db, watch};
+use ty_project::{CollectReporter, Db, watch};
 use ty_project::{ProjectDatabase, ProjectMetadata};
 use ty_server::run_server;
 
@@ -275,7 +276,8 @@ impl MainLoop {
                     rayon::spawn(move || {
                         match salsa::Cancelled::catch(|| {
                             let mut reporter = IndicatifReporter::from(self.printer);
-                            db.check_with_reporter(&mut reporter)
+                            db.check_with_reporter(&mut reporter);
+                            reporter.collector.into_sorted(&db)
                         }) {
                             Ok(result) => {
                                 // Send the result back to the main loop for printing.
@@ -394,8 +396,12 @@ impl MainLoop {
     }
 }
 
+struct IndicatifReporter {
+    collector: CollectReporter,
+    state: IndicatifReporterState,
+}
 /// A progress reporter for `ty check`.
-enum IndicatifReporter {
+enum IndicatifReporterState {
     /// A constructed reporter that is not yet ready, contains the target for the progress bar.
     Pending(indicatif::ProgressDrawTarget),
     /// A reporter that is ready, containing a progress bar to report to.
@@ -408,18 +414,25 @@ enum IndicatifReporter {
 
 impl From<Printer> for IndicatifReporter {
     fn from(printer: Printer) -> Self {
-        Self::Pending(printer.progress_target())
+        Self {
+            collector: CollectReporter::default(),
+            state: IndicatifReporterState::Pending(printer.progress_target()),
+        }
     }
 }
 
 impl ty_project::ProgressReporter for IndicatifReporter {
     fn set_files(&mut self, files: usize) {
+        self.collector.set_files(files);
+
         let target = match std::mem::replace(
-            self,
-            IndicatifReporter::Pending(indicatif::ProgressDrawTarget::hidden()),
+            &mut self.state,
+            IndicatifReporterState::Pending(indicatif::ProgressDrawTarget::hidden()),
         ) {
-            Self::Pending(target) => target,
-            Self::Initialized(_) => panic!("The progress reporter should only be initialized once"),
+            IndicatifReporterState::Pending(target) => target,
+            IndicatifReporterState::Initialized(_) => {
+                panic!("The progress reporter should only be initialized once")
+            }
         };
 
         let bar = indicatif::ProgressBar::with_draw_target(Some(files as u64), target);
@@ -431,19 +444,24 @@ impl ty_project::ProgressReporter for IndicatifReporter {
             .progress_chars("--"),
         );
         bar.set_message("Checking");
-        *self = Self::Initialized(bar);
+        self.state = IndicatifReporterState::Initialized(bar);
     }
 
-    fn report_file(&self, _file: &ruff_db::files::File) {
-        match self {
-            IndicatifReporter::Initialized(progress_bar) => {
+    fn report_file(&self, db: &dyn Db, file: File, diagnostics: &[Diagnostic]) {
+        self.collector.report_file(db, file, diagnostics);
+        match &self.state {
+            IndicatifReporterState::Initialized(progress_bar) => {
                 progress_bar.inc(1);
             }
-            IndicatifReporter::Pending(_) => {
+            IndicatifReporterState::Pending(_) => {
                 panic!("`report_file` called before `set_files`")
             }
         }
     }
+
+    fn report_diagnostics(&mut self, db: &dyn Db, diagnostics: Vec<Diagnostic>) {
+        self.collector.report_diagnostics(db, diagnostics);
+    }
 }
 
 #[derive(Debug)]

crates/ty/tests/cli/rule_selection.rs

@@ -863,10 +863,18 @@ fn overrides_unknown_rules() -> anyhow::Result<()> {
         ),
     ])?;
 
-    assert_cmd_snapshot!(case.command(), @r#"
+    assert_cmd_snapshot!(case.command(), @r###"
     success: false
     exit_code: 1
     ----- stdout -----
+    error[division-by-zero]: Cannot divide object of type `Literal[4]` by zero
+     --> main.py:2:5
+      |
+    2 | y = 4 / 0
+      |     ^^^^^
+      |
+    info: rule `division-by-zero` was selected in the configuration file
+
     warning[unknown-rule]: Unknown lint rule `division-by-zer`
       --> pyproject.toml:10:1
        |
@@ -876,14 +884,6 @@ fn overrides_unknown_rules() -> anyhow::Result<()> {
        | ^^^^^^^^^^^^^^^
        |
 
-    error[division-by-zero]: Cannot divide object of type `Literal[4]` by zero
-     --> main.py:2:5
-      |
-    2 | y = 4 / 0
-      |     ^^^^^
-      |
-    info: rule `division-by-zero` was selected in the configuration file
-
     warning[division-by-zero]: Cannot divide object of type `Literal[4]` by zero
      --> tests/test_main.py:2:5
       |
@@ -896,7 +896,7 @@ fn overrides_unknown_rules() -> anyhow::Result<()> {
 
     ----- stderr -----
     WARN ty is pre-release software and not ready for production use. Expect to encounter bugs, missing features, and fatal errors.
-    "#);
+    "###);
 
     Ok(())
 }

crates/ty_project/src/db.rs

@@ -5,7 +5,7 @@ use std::{cmp, fmt};
 
 pub use self::changes::ChangeResult;
 use crate::metadata::settings::file_settings;
-use crate::{DEFAULT_LINT_REGISTRY, DummyReporter};
+use crate::{CollectReporter, DEFAULT_LINT_REGISTRY};
 use crate::{ProgressReporter, Project, ProjectMetadata};
 use ruff_db::Db as SourceDb;
 use ruff_db::diagnostic::Diagnostic;
@@ -87,16 +87,18 @@ impl ProjectDatabase {
     ///
     /// [`set_check_mode`]: ProjectDatabase::set_check_mode
     pub fn check(&self) -> Vec<Diagnostic> {
-        self.project().check(self, &mut DummyReporter)
+        let mut collector = CollectReporter::default();
+        self.project().check(self, &mut collector);
+        collector.into_sorted(self)
     }
 
     /// Checks the files in the project and its dependencies, using the given reporter.
     ///
     /// Use [`set_check_mode`] to update the check mode.
     ///
     /// [`set_check_mode`]: ProjectDatabase::set_check_mode
-    pub fn check_with_reporter(&self, reporter: &mut dyn ProgressReporter) -> Vec<Diagnostic> {
-        self.project().check(self, reporter)
+    pub fn check_with_reporter(&self, reporter: &mut dyn ProgressReporter) {
+        self.project().check(self, reporter);
     }
 
     #[tracing::instrument(level = "debug", skip(self))]

crates/ty_project/src/lib.rs

@@ -128,16 +128,45 @@ pub trait ProgressReporter: Send + Sync {
     fn set_files(&mut self, files: usize);
 
     /// Report the completion of a given file.
-    fn report_file(&self, file: &File);
+    fn report_file(&self, db: &dyn Db, file: File, diagnostics: &[Diagnostic]);
+
+    /// Reports settings or IO related diagnostics. The diagnostics
+    /// can belong to different files or no file at all.
+    /// But it's never a file for which [`Self::report_file`] gets called.
+    fn report_diagnostics(&mut self, db: &dyn Db, diagnostics: Vec<Diagnostic>);
 }
 
-/// A no-op implementation of [`ProgressReporter`].
+/// Reporter that collects all diagnostics into a `Vec`.
 #[derive(Default)]
-pub struct DummyReporter;
+pub struct CollectReporter(std::sync::Mutex<Vec<Diagnostic>>);
+
+impl CollectReporter {
+    pub fn into_sorted(self, db: &dyn Db) -> Vec<Diagnostic> {
+        let mut diagnostics = self.0.into_inner().unwrap();
+        diagnostics.sort_by(|left, right| {
+            left.rendering_sort_key(db)
+                .cmp(&right.rendering_sort_key(db))
+        });
+        diagnostics
+    }
+}
 
-impl ProgressReporter for DummyReporter {
+impl ProgressReporter for CollectReporter {
     fn set_files(&mut self, _files: usize) {}
-    fn report_file(&self, _file: &File) {}
+    fn report_file(&self, _db: &dyn Db, _file: File, diagnostics: &[Diagnostic]) {
+        if diagnostics.is_empty() {
+            return;
+        }
+
+        self.0
+            .lock()
+            .unwrap()
+            .extend(diagnostics.iter().map(Clone::clone));
+    }
+
+    fn report_diagnostics(&mut self, _db: &dyn Db, diagnostics: Vec<Diagnostic>) {
+        self.0.get_mut().unwrap().extend(diagnostics);
+    }
 }
 
 #[salsa::tracked]
@@ -225,11 +254,7 @@ impl Project {
     }
 
     /// Checks the project and its dependencies according to the project's check mode.
-    pub(crate) fn check(
-        self,
-        db: &ProjectDatabase,
-        reporter: &mut dyn ProgressReporter,
-    ) -> Vec<Diagnostic> {
+    pub(crate) fn check(self, db: &ProjectDatabase, reporter: &mut dyn ProgressReporter) {
         let project_span = tracing::debug_span!("Project::check");
         let _span = project_span.enter();
 
@@ -239,12 +264,11 @@ impl Project {
             name = self.name(db)
         );
 
-        let mut diagnostics: Vec<Diagnostic> = Vec::new();
-        diagnostics.extend(
-            self.settings_diagnostics(db)
-                .iter()
-                .map(OptionDiagnostic::to_diagnostic),
-        );
+        let mut diagnostics: Vec<Diagnostic> = self
+            .settings_diagnostics(db)
+            .iter()
+            .map(OptionDiagnostic::to_diagnostic)
+            .collect();
 
         let files = ProjectFiles::new(db, self);
         reporter.set_files(files.len());
@@ -256,30 +280,27 @@ impl Project {
                 .map(IOErrorDiagnostic::to_diagnostic),
         );
 
+        reporter.report_diagnostics(db, diagnostics);
+
         let open_files = self.open_files(db);
         let check_start = ruff_db::Instant::now();
-        let file_diagnostics = std::sync::Mutex::new(vec![]);
 
         {
             let db = db.clone();
-            let file_diagnostics = &file_diagnostics;
             let project_span = &project_span;
-            let reporter = &reporter;
 
             rayon::scope(move |scope| {
                 for file in &files {
                     let db = db.clone();
+                    let reporter = &*reporter;
                     scope.spawn(move |_| {
                         let check_file_span =
                             tracing::debug_span!(parent: project_span, "check_file", ?file);
                         let _entered = check_file_span.entered();
 
                         match check_file_impl(&db, file) {
                             Ok(diagnostics) => {
-                                file_diagnostics
-                                    .lock()
-                                    .unwrap()
-                                    .extend(diagnostics.iter().map(Clone::clone));
+                                reporter.report_file(&db, file, diagnostics);
 
                                 // This is outside `check_file_impl` to avoid that opening or closing
                                 // a file invalidates the `check_file_impl` query of every file!
@@ -295,28 +316,18 @@ impl Project {
                                 }
                             }
                             Err(io_error) => {
-                                file_diagnostics.lock().unwrap().push(io_error.clone());
+                                reporter.report_file(&db, file, std::slice::from_ref(io_error));
                             }
                         }
-
-                        reporter.report_file(&file);
                     });
                 }
             });
-        }
+        };
 
         tracing::debug!(
             "Checking all files took {:.3}s",
             check_start.elapsed().as_secs_f64(),
         );
-
-        let mut file_diagnostics = file_diagnostics.into_inner().unwrap();
-        file_diagnostics.sort_by(|left, right| {
-            left.rendering_sort_key(db)
-                .cmp(&right.rendering_sort_key(db))
-        });
-        diagnostics.extend(file_diagnostics);
-        diagnostics
     }
 
     pub(crate) fn check_file(self, db: &dyn Db, file: File) -> Vec<Diagnostic> {

crates/ty_server/src/lib.rs

@@ -5,7 +5,7 @@ use lsp_server::Connection;
 use ruff_db::system::{OsSystem, SystemPathBuf};
 
 pub use crate::logging::{LogLevel, init_logging};
-pub use crate::server::Server;
+pub use crate::server::{PartialWorkspaceProgress, PartialWorkspaceProgressParams, Server};
 pub use crate::session::{ClientOptions, DiagnosticMode};
 pub use document::{NotebookDocument, PositionEncoding, TextDocument};
 pub(crate) use session::{DocumentQuery, Session};

crates/ty_server/src/server.rs

@@ -29,6 +29,7 @@ pub(crate) use main_loop::{
     Action, ConnectionSender, Event, MainLoopReceiver, MainLoopSender, SendRequest,
 };
 pub(crate) type Result<T> = std::result::Result<T, api::Error>;
+pub use api::{PartialWorkspaceProgress, PartialWorkspaceProgressParams};
 
 pub struct Server {
     connection: Connection,

crates/ty_server/src/server/api.rs

@@ -19,6 +19,7 @@ use self::traits::{NotificationHandler, RequestHandler};
 use super::{Result, schedule::BackgroundSchedule};
 use crate::session::client::Client;
 pub(crate) use diagnostics::publish_settings_diagnostics;
+pub use requests::{PartialWorkspaceProgress, PartialWorkspaceProgressParams};
 use ruff_db::panic::PanicError;
 
 /// Processes a request from the client to the server.

crates/ty_server/src/server/api/requests.rs

@@ -33,3 +33,5 @@ pub(super) use shutdown::ShutdownHandler;
 pub(super) use signature_help::SignatureHelpRequestHandler;
 pub(super) use workspace_diagnostic::WorkspaceDiagnosticRequestHandler;
 pub(super) use workspace_symbols::WorkspaceSymbolRequestHandler;
+
+pub use workspace_diagnostic::{PartialWorkspaceProgress, PartialWorkspaceProgressParams};