[red-knot] Allow module-resolution options to be specified via the CLI (#12246)

Cargo.lock

@@ -1864,6 +1864,7 @@ name = "red_knot"
 version = "0.0.0"
 dependencies = [
  "anyhow",
+ "clap",
  "countme",
  "crossbeam",
  "ctrlc",

crates/red_knot/Cargo.toml

@@ -19,6 +19,7 @@ ruff_db = { workspace = true }
 ruff_python_ast = { workspace = true }
 
 anyhow = { workspace = true }
+clap = { workspace = true, features = ["wrap_help"] }
 countme = { workspace = true, features = ["enable"] }
 crossbeam = { workspace = true }
 ctrlc = { version = "3.4.4" }

crates/red_knot/src/lib.rs

@@ -8,6 +8,7 @@ use crate::db::Jar;
 pub mod db;
 pub mod lint;
 pub mod program;
+pub mod target_version;
 pub mod watch;
 
 #[derive(Debug, Clone)]

crates/red_knot/src/main.rs

@@ -1,5 +1,6 @@
 use std::sync::Mutex;
 
+use clap::Parser;
 use crossbeam::channel as crossbeam_channel;
 use salsa::ParallelDatabase;
 use tracing::subscriber::Interest;
@@ -10,13 +11,38 @@ use tracing_subscriber::{Layer, Registry};
 use tracing_tree::time::Uptime;
 
 use red_knot::program::{FileWatcherChange, Program};
+use red_knot::target_version::TargetVersion;
 use red_knot::watch::FileWatcher;
 use red_knot::Workspace;
-use red_knot_module_resolver::{
-    set_module_resolution_settings, RawModuleResolutionSettings, TargetVersion,
-};
+use red_knot_module_resolver::{set_module_resolution_settings, RawModuleResolutionSettings};
 use ruff_db::files::system_path_to_file;
-use ruff_db::system::{OsSystem, System, SystemPath};
+use ruff_db::system::{OsSystem, System, SystemPath, SystemPathBuf};
+
+#[derive(Debug, Parser)]
+#[command(
+    author,
+    name = "red-knot",
+    about = "An experimental multifile analysis backend for Ruff"
+)]
+#[command(version)]
+struct Args {
+    #[clap(help = "File to check", required = true, value_name = "FILE")]
+    entry_point: SystemPathBuf,
+    #[arg(
+        long,
+        value_name = "DIRECTORY",
+        help = "Custom directory to use for stdlib typeshed stubs"
+    )]
+    custom_typeshed_dir: Option<SystemPathBuf>,
+    #[arg(
+        long,
+        value_name = "PATH",
+        help = "Additional path to use as a module-resolution source (can be passed multiple times)"
+    )]
+    extra_search_path: Vec<SystemPathBuf>,
+    #[arg(long, help = "Python version to assume when resolving types", default_value_t = TargetVersion::default(), value_name="VERSION")]
+    target_version: TargetVersion,
+}
 
 #[allow(
     clippy::print_stdout,
@@ -28,30 +54,35 @@ pub fn main() -> anyhow::Result<()> {
     countme::enable(true);
     setup_tracing();
 
-    let arguments: Vec<_> = std::env::args().collect();
+    let Args {
+        entry_point,
+        custom_typeshed_dir,
+        extra_search_path: extra_search_paths,
+        target_version,
+    } = Args::parse_from(std::env::args().collect::<Vec<_>>());
 
-    if arguments.len() < 2 {
-        eprintln!("Usage: red_knot <path>");
-        return Err(anyhow::anyhow!("Invalid arguments"));
+    tracing::trace!("Target version: {target_version}");
+    if let Some(custom_typeshed) = custom_typeshed_dir.as_ref() {
+        tracing::trace!("Custom typeshed directory: {custom_typeshed}");
+    }
+    if !extra_search_paths.is_empty() {
+        tracing::trace!("extra search paths: {extra_search_paths:?}");
     }
 
     let cwd = std::env::current_dir().unwrap();
     let cwd = SystemPath::from_std_path(&cwd).unwrap();
     let system = OsSystem::new(cwd);
-    let entry_point = SystemPath::new(&arguments[1]);
 
-    if !system.path_exists(entry_point) {
+    if !system.path_exists(&entry_point) {
         eprintln!("The entry point does not exist.");
         return Err(anyhow::anyhow!("Invalid arguments"));
     }
 
-    if !system.is_file(entry_point) {
+    if !system.is_file(&entry_point) {
         eprintln!("The entry point is not a file.");
         return Err(anyhow::anyhow!("Invalid arguments"));
     }
 
-    let entry_point = entry_point.to_path_buf();
-
     let workspace_folder = entry_point.parent().unwrap();
     let workspace = Workspace::new(workspace_folder.to_path_buf());
 
@@ -62,11 +93,11 @@ pub fn main() -> anyhow::Result<()> {
     set_module_resolution_settings(
         &mut program,
         RawModuleResolutionSettings {
-            extra_paths: vec![],
+            extra_paths: extra_search_paths,
             workspace_root: workspace_search_path,
             site_packages: None,
-            custom_typeshed: None,
-            target_version: TargetVersion::Py38,
+            custom_typeshed: custom_typeshed_dir,
+            target_version: red_knot_module_resolver::TargetVersion::from(target_version),
         },
     );
 

crates/red_knot/src/target_version.rs

@@ -0,0 +1,50 @@
+use std::fmt;
+
+/// Enumeration of all supported Python versions
+///
+/// TODO: unify with the `PythonVersion` enum in the linter/formatter crates?
+#[derive(Copy, Clone, Hash, Debug, PartialEq, Eq, PartialOrd, Ord, Default, clap::ValueEnum)]
+pub enum TargetVersion {
+    Py37,
+    #[default]
+    Py38,
+    Py39,
+    Py310,
+    Py311,
+    Py312,
+    Py313,
+}
+
+impl TargetVersion {
+    const fn as_str(self) -> &'static str {
+        match self {
+            Self::Py37 => "py37",
+            Self::Py38 => "py38",
+            Self::Py39 => "py39",
+            Self::Py310 => "py310",
+            Self::Py311 => "py311",
+            Self::Py312 => "py312",
+            Self::Py313 => "py313",
+        }
+    }
+}
+
+impl fmt::Display for TargetVersion {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(self.as_str())
+    }
+}
+
+impl From<TargetVersion> for red_knot_module_resolver::TargetVersion {
+    fn from(value: TargetVersion) -> Self {
+        match value {
+            TargetVersion::Py37 => red_knot_module_resolver::TargetVersion::Py37,
+            TargetVersion::Py38 => red_knot_module_resolver::TargetVersion::Py38,
+            TargetVersion::Py39 => red_knot_module_resolver::TargetVersion::Py39,
+            TargetVersion::Py310 => red_knot_module_resolver::TargetVersion::Py310,
+            TargetVersion::Py311 => red_knot_module_resolver::TargetVersion::Py311,
+            TargetVersion::Py312 => red_knot_module_resolver::TargetVersion::Py312,
+            TargetVersion::Py313 => red_knot_module_resolver::TargetVersion::Py313,
+        }
+    }
+}