Gracefully handle errors in CLI (#12747)

crates/red_knot/src/main.rs

@@ -1,9 +1,11 @@
-use std::process::ExitCode;
+use std::process::{ExitCode, Termination};
 use std::sync::Mutex;
 
+use anyhow::{anyhow, Context};
 use clap::Parser;
 use colored::Colorize;
 use crossbeam::channel as crossbeam_channel;
+use salsa::plumbing::ZalsaDatabase;
 
 use red_knot_server::run_server;
 use red_knot_workspace::db::RootDatabase;
@@ -12,7 +14,7 @@ use red_knot_workspace::watch;
 use red_knot_workspace::watch::WorkspaceWatcher;
 use red_knot_workspace::workspace::WorkspaceMetadata;
 use ruff_db::program::{ProgramSettings, SearchPathSettings};
-use ruff_db::system::{OsSystem, System, SystemPathBuf};
+use ruff_db::system::{OsSystem, System, SystemPath, SystemPathBuf};
 use target_version::TargetVersion;
 
 use crate::logging::{setup_tracing, Verbosity};
@@ -86,30 +88,25 @@ pub enum Command {
 }
 
 #[allow(clippy::print_stdout, clippy::unnecessary_wraps, clippy::print_stderr)]
-pub fn main() -> ExitCode {
-    match run() {
-        Ok(status) => status.into(),
-        Err(error) => {
-            {
-                use std::io::Write;
-
-                // Use `writeln` instead of `eprintln` to avoid panicking when the stderr pipe is broken.
-                let mut stderr = std::io::stderr().lock();
-
-                // This communicates that this isn't a linter error but ruff itself hard-errored for
-                // some reason (e.g. failed to resolve the configuration)
-                writeln!(stderr, "{}", "ruff failed".red().bold()).ok();
-                // Currently we generally only see one error, but e.g. with io errors when resolving
-                // the configuration it is help to chain errors ("resolving configuration failed" ->
-                // "failed to read file: subdir/pyproject.toml")
-                for cause in error.chain() {
-                    writeln!(stderr, "  {} {cause}", "Cause:".bold()).ok();
-                }
-            }
-
-            ExitStatus::Error.into()
+pub fn main() -> ExitStatus {
+    run().unwrap_or_else(|error| {
+        use std::io::Write;
+
+        // Use `writeln` instead of `eprintln` to avoid panicking when the stderr pipe is broken.
+        let mut stderr = std::io::stderr().lock();
+
+        // This communicates that this isn't a linter error but Red Knot itself hard-errored for
+        // some reason (e.g. failed to resolve the configuration)
+        writeln!(stderr, "{}", "ruff failed".red().bold()).ok();
+        // Currently we generally only see one error, but e.g. with io errors when resolving
+        // the configuration it is help to chain errors ("resolving configuration failed" ->
+        // "failed to read file: subdir/pyproject.toml")
+        for cause in error.chain() {
+            writeln!(stderr, "  {} {cause}", "Cause:".bold()).ok();
         }
-    }
+
+        ExitStatus::Error
+    })
 }
 
 fn run() -> anyhow::Result<ExitStatus> {
@@ -132,28 +129,43 @@ fn run() -> anyhow::Result<ExitStatus> {
     countme::enable(verbosity.is_trace());
     let _guard = setup_tracing(verbosity)?;
 
-    let cwd = if let Some(cwd) = current_directory {
-        let canonicalized = cwd.as_utf8_path().canonicalize_utf8().unwrap();
-        SystemPathBuf::from_utf8_path_buf(canonicalized)
-    } else {
-        let cwd = std::env::current_dir().unwrap();
-        SystemPathBuf::from_path_buf(cwd).unwrap()
+    // The base path to which all CLI arguments are relative to.
+    let cli_base_path = {
+        let cwd = std::env::current_dir().context("Failed to get the current working directory")?;
+        SystemPathBuf::from_path_buf(cwd).map_err(|path| anyhow!("The current working directory '{}' contains non-unicode characters. Red Knot only supports unicode paths.", path.display()))?
     };
 
+    let cwd = current_directory
+        .map(|cwd| {
+            if cwd.as_std_path().is_dir() {
+                Ok(SystemPath::absolute(&cwd, &cli_base_path))
+            } else {
+                Err(anyhow!(
+                    "Provided current-directory path '{cwd}' is not a directory."
+                ))
+            }
+        })
+        .transpose()?
+        .unwrap_or_else(|| cli_base_path.clone());
+
     let system = OsSystem::new(cwd.clone());
-    let workspace_metadata =
-        WorkspaceMetadata::from_path(system.current_directory(), &system).unwrap();
-
-    let site_packages = if let Some(venv_path) = venv_path {
-        let venv_path = system.canonicalize_path(&venv_path).unwrap_or(venv_path);
-        assert!(
-            system.is_directory(&venv_path),
-            "Provided venv-path {venv_path} is not a directory!"
-        );
-        site_packages_dirs_of_venv(&venv_path, &system).unwrap()
-    } else {
-        vec![]
-    };
+    let workspace_metadata = WorkspaceMetadata::from_path(system.current_directory(), &system)?;
+
+    // TODO: Verify the remaining search path settings eagerly.
+    let site_packages = venv_path
+        .map(|venv_path| {
+            let venv_path = SystemPath::absolute(venv_path, &cli_base_path);
+
+            if system.is_directory(&venv_path) {
+                Ok(site_packages_dirs_of_venv(&venv_path, &system)?)
+            } else {
+                Err(anyhow!(
+                    "Provided venv-path {venv_path} is not a directory!"
+                ))
+            }
+        })
+        .transpose()?
+        .unwrap_or_default();
 
     // TODO: Respect the settings from the workspace metadata. when resolving the program settings.
     let program_settings = ProgramSettings {
@@ -207,9 +219,9 @@ pub enum ExitStatus {
     Error = 2,
 }
 
-impl From<ExitStatus> for ExitCode {
-    fn from(status: ExitStatus) -> Self {
-        ExitCode::from(status as u8)
+impl Termination for ExitStatus {
+    fn report(self) -> ExitCode {
+        ExitCode::from(self as u8)
     }
 }
 
@@ -262,12 +274,11 @@ impl MainLoop {
         result
     }
 
-    #[allow(clippy::print_stderr)]
     fn main_loop(&mut self, db: &mut RootDatabase) -> ExitStatus {
         // Schedule the first check.
         tracing::debug!("Starting main loop");
 
-        let mut revision = 0usize;
+        let mut revision = 0u64;
 
         while let Ok(message) = self.receiver.recv() {
             match message {
@@ -282,7 +293,7 @@ impl MainLoop {
                             // Send the result back to the main loop for printing.
                             sender
                                 .send(MainLoopMessage::CheckCompleted { result, revision })
-                                .ok();
+                                .unwrap();
                         }
                     });
                 }
@@ -291,17 +302,20 @@ impl MainLoop {
                     result,
                     revision: check_revision,
                 } => {
+                    let has_diagnostics = !result.is_empty();
                     if check_revision == revision {
-                        eprintln!("{}", result.join("\n"));
+                        for diagnostic in result {
+                            tracing::error!("{}", diagnostic);
+                        }
                     } else {
                         tracing::debug!("Discarding check result for outdated revision: current: {revision}, result revision: {check_revision}");
                     }
 
                     if self.watcher.is_none() {
-                        return if result.is_empty() {
-                            ExitStatus::Success
-                        } else {
+                        return if has_diagnostics {
                             ExitStatus::Failure
+                        } else {
+                            ExitStatus::Success
                         };
                     }
 
@@ -318,6 +332,10 @@ impl MainLoop {
                     self.sender.send(MainLoopMessage::CheckWorkspace).unwrap();
                 }
                 MainLoopMessage::Exit => {
+                    // Cancel any pending queries and wait for them to complete.
+                    // TODO: Don't use Salsa internal APIs
+                    //  [Zulip-Thread](https://salsa.zulipchat.com/#narrow/stream/333573-salsa-3.2E0/topic/Expose.20an.20API.20to.20cancel.20other.20queries)
+                    let _ = db.zalsa_mut();
                     return ExitStatus::Success;
                 }
             }
@@ -344,10 +362,7 @@ impl MainLoopCancellationToken {
 #[derive(Debug)]
 enum MainLoopMessage {
     CheckWorkspace,
-    CheckCompleted {
-        result: Vec<String>,
-        revision: usize,
-    },
+    CheckCompleted { result: Vec<String>, revision: u64 },
     ApplyChanges(Vec<watch::ChangeEvent>),
     Exit,
 }

crates/red_knot_workspace/src/site_packages.rs

@@ -34,12 +34,17 @@ fn site_packages_dir_from_sys_prefix(
     sys_prefix_path: &SystemPath,
     system: &dyn System,
 ) -> Result<SystemPathBuf, SitePackagesDiscoveryError> {
+    tracing::debug!("Searching for site-packages directory in '{sys_prefix_path}'");
+
     if cfg!(target_os = "windows") {
         let site_packages = sys_prefix_path.join("Lib/site-packages");
-        return system
-            .is_directory(&site_packages)
-            .then_some(site_packages)
-            .ok_or(SitePackagesDiscoveryError::NoSitePackagesDirFound);
+
+        return if system.is_directory(&site_packages) {
+            tracing::debug!("Resolved site-packages directory to '{site_packages}'");
+            Ok(site_packages)
+        } else {
+            Err(SitePackagesDiscoveryError::NoSitePackagesDirFound)
+        };
     }
 
     // In the Python standard library's `site.py` module (used for finding `site-packages`
@@ -68,11 +73,12 @@ fn site_packages_dir_from_sys_prefix(
         let Ok(entry) = entry_result else {
             continue;
         };
+
         if !entry.file_type().is_directory() {
             continue;
         }
 
-        let path = entry.path();
+        let mut path = entry.into_path();
 
         // The `python3.x` part of the `site-packages` path can't be computed from
         // the `--target-version` the user has passed, as they might be running Python 3.12 locally
@@ -84,29 +90,26 @@ fn site_packages_dir_from_sys_prefix(
         // the `pyvenv.cfg` file anyway, in which case we could switch to that method
         // rather than iterating through the whole directory until we find
         // an entry where the last component of the path starts with `python3.`
-        if !path
-            .components()
-            .next_back()
-            .is_some_and(|last_part| last_part.as_str().starts_with("python3."))
-        {
+        let name = path
+            .file_name()
+            .expect("File name to be non-null because path is guaranteed to be a child of `lib`");
+
+        if !name.starts_with("python3.") {
             continue;
         }
 
-        let site_packages_candidate = path.join("site-packages");
-        if system.is_directory(&site_packages_candidate) {
-            tracing::debug!(
-                "Resoled site-packages directory: {}",
-                site_packages_candidate
-            );
-            return Ok(site_packages_candidate);
+        path.push("site-packages");
+        if system.is_directory(&path) {
+            tracing::debug!("Resolved site-packages directory to '{path}'");
+            return Ok(path);
         }
     }
     Err(SitePackagesDiscoveryError::NoSitePackagesDirFound)
 }
 
 #[derive(Debug, thiserror::Error)]
 pub enum SitePackagesDiscoveryError {
-    #[error("Failed to search the virtual environment directory for `site-packages` due to {0}")]
+    #[error("Failed to search the virtual environment directory for `site-packages`")]
     CouldNotReadLibDirectory(#[from] io::Error),
     #[error("Could not find the `site-packages` directory in the virtual environment")]
     NoSitePackagesDirFound,

crates/red_knot_workspace/src/workspace/metadata.rs

@@ -22,15 +22,13 @@ pub struct PackageMetadata {
 impl WorkspaceMetadata {
     /// Discovers the closest workspace at `path` and returns its metadata.
     pub fn from_path(path: &SystemPath, system: &dyn System) -> anyhow::Result<WorkspaceMetadata> {
-        let root = if system.is_file(path) {
-            path.parent().unwrap().to_path_buf()
-        } else {
-            path.to_path_buf()
-        };
+        assert!(
+            system.is_directory(path),
+            "Workspace root path must be a directory"
+        );
+        tracing::debug!("Searching for workspace in '{path}'");
 
-        if !system.is_directory(&root) {
-            anyhow::bail!("no workspace found at {:?}", root);
-        }
+        let root = path.to_path_buf();
 
         // TODO: Discover package name from `pyproject.toml`.
         let package_name: Name = path.file_name().unwrap_or("<root>").into();