Allow passing a virtual environment to ruff analyze graph

crates/ruff/src/args.rs

@@ -177,6 +177,9 @@ pub struct AnalyzeGraphCommand {
     /// The minimum Python version that should be supported.
     #[arg(long, value_enum)]
     target_version: Option<PythonVersion>,
+    /// Path to a virtual environment to use for resolving additional dependencies
+    #[arg(long)]
+    python: Option<PathBuf>,
 }
 
 // The `Parser` derive is for ruff_dev, for ruff `Args` would be sufficient
@@ -796,6 +799,7 @@ impl AnalyzeGraphCommand {
         let format_arguments = AnalyzeGraphArgs {
             files: self.files,
             direction: self.direction,
+            python: self.python,
         };
 
         let cli_overrides = ExplicitConfigOverrides {
@@ -1261,6 +1265,7 @@ impl LineColumnParseError {
 pub struct AnalyzeGraphArgs {
     pub files: Vec<PathBuf>,
     pub direction: Direction,
+    pub python: Option<PathBuf>,
 }
 
 /// Configuration overrides provided via dedicated CLI flags:

crates/ruff/src/commands/analyze_graph.rs

@@ -75,6 +75,8 @@ pub(crate) fn analyze_graph(
             .target_version
             .as_tuple()
             .into(),
+        args.python
+            .and_then(|python| SystemPathBuf::from_path_buf(python).ok()),
     )?;
 
     let imports = {

crates/ruff/tests/analyze_graph.rs

@@ -422,3 +422,153 @@ fn nested_imports() -> Result<()> {
 
     Ok(())
 }
+
+/// Test for venv resolution with the `--python` flag.
+///
+/// Based on the [albatross-virtual-workspace] example from the uv repo and the report in [#16598].
+///
+/// [albatross-virtual-workspace]: https://github.com/astral-sh/uv/tree/aa629c4a/scripts/workspaces/albatross-virtual-workspace
+/// [#16598]: https://github.com/astral-sh/ruff/issues/16598
+#[test]
+fn venv() -> Result<()> {
+    let tempdir = TempDir::new()?;
+    let root = ChildPath::new(tempdir.path());
+
+    // packages
+    // ├── albatross
+    // │   ├── check_installed_albatross.py
+    // │   ├── pyproject.toml
+    // │   └── src
+    // │       └── albatross
+    // │           └── __init__.py
+    // └── bird-feeder
+    //     ├── check_installed_bird_feeder.py
+    //     ├── pyproject.toml
+    //     └── src
+    //         └── bird_feeder
+    //             └── __init__.py
+
+    let packages = root.child("packages");
+
+    let albatross = packages.child("albatross");
+    albatross
+        .child("check_installed_albatross.py")
+        .write_str("from albatross import fly")?;
+    albatross
+        .child("pyproject.toml")
+        .write_str(indoc::indoc! {r#"
+        [project]
+        name = "albatross"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = ["bird-feeder", "tqdm>=4,<5"]
+
+        [tool.uv.sources]
+        bird-feeder = { workspace = true }
+    "#})?;
+    albatross
+        .child("src")
+        .child("albatross")
+        .child("__init__.py")
+        .write_str("import tqdm; from bird_feeder import use")?;
+
+    let bird_feeder = packages.child("bird-feeder");
+    bird_feeder
+        .child("check_installed_bird_feeder.py")
+        .write_str("from bird_feeder import use; from albatross import fly")?;
+    bird_feeder
+        .child("pyproject.toml")
+        .write_str(indoc::indoc! {r#"
+        [project]
+        name = "bird-feeder"
+        version = "1.0.0"
+        requires-python = ">=3.12"
+        dependencies = ["anyio>=4.3.0,<5"]
+    "#})?;
+    bird_feeder
+        .child("src")
+        .child("bird_feeder")
+        .child("__init__.py")
+        .write_str("import anyio")?;
+
+    let venv = root.child(".venv");
+    let bin = venv.child("bin");
+    bin.child("python").touch()?;
+    let home = format!("home = {}", bin.to_string_lossy());
+    venv.child("pyvenv.cfg").write_str(&home)?;
+    let site_packages = venv.child("lib").child("python3.12").child("site-packages");
+    site_packages
+        .child("_albatross.pth")
+        .write_str(&albatross.join("src").to_string_lossy())?;
+    site_packages
+        .child("_bird_feeder.pth")
+        .write_str(&bird_feeder.join("src").to_string_lossy())?;
+    site_packages.child("tqdm").child("__init__.py").touch()?;
+
+    // without `--python .venv`, the result should only include dependencies within the albatross
+    // package
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec(),
+    }, {
+        assert_cmd_snapshot!(
+            command().arg("packages/albatross").current_dir(&root),
+            @r#"
+        success: true
+        exit_code: 0
+        ----- stdout -----
+        {
+          "packages/albatross/check_installed_albatross.py": [
+            "packages/albatross/src/albatross/__init__.py"
+          ],
+          "packages/albatross/src/albatross/__init__.py": []
+        }
+
+        ----- stderr -----
+        "#);
+    });
+
+    // with `--python .venv` both workspace and third-party dependencies are included
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec(),
+    }, {
+        assert_cmd_snapshot!(
+            command().args(["--python", ".venv"]).arg("packages/albatross").current_dir(&root),
+            @r#"
+        success: true
+        exit_code: 0
+        ----- stdout -----
+        {
+          "packages/albatross/check_installed_albatross.py": [
+            "packages/albatross/src/albatross/__init__.py"
+          ],
+          "packages/albatross/src/albatross/__init__.py": [
+            ".venv/lib/python3.12/site-packages/tqdm/__init__.py",
+            "packages/bird-feeder/src/bird_feeder/__init__.py"
+          ]
+        }
+
+        ----- stderr -----
+        "#);
+    });
+
+    // test the error message for a non-existent venv. it's important that the `ruff analyze graph`
+    // flag matches the red-knot flag used to generate the error message (`--python`)
+    insta::with_settings!({
+        filters => INSTA_FILTERS.to_vec(),
+    }, {
+        assert_cmd_snapshot!(
+            command().args(["--python", "none"]).arg("packages/albatross").current_dir(&root),
+            @r"
+        success: false
+        exit_code: 2
+        ----- stdout -----
+
+        ----- stderr -----
+        ruff failed
+          Cause: Invalid search path settings
+          Cause: Failed to discover the site-packages directory: Invalid `--python` argument: `none` could not be canonicalized
+        ");
+    });
+
+    Ok(())
+}

crates/ruff_graph/src/db.rs

@@ -4,7 +4,8 @@ use zip::CompressionMethod;
 
 use red_knot_python_semantic::lint::{LintRegistry, RuleSelection};
 use red_knot_python_semantic::{
-    default_lint_registry, Db, Program, ProgramSettings, PythonPlatform, SearchPathSettings,
+    default_lint_registry, Db, Program, ProgramSettings, PythonPath, PythonPlatform,
+    SearchPathSettings,
 };
 use ruff_db::files::{File, Files};
 use ruff_db::system::{OsSystem, System, SystemPathBuf};
@@ -32,8 +33,12 @@ impl ModuleDb {
     pub fn from_src_roots(
         src_roots: Vec<SystemPathBuf>,
         python_version: PythonVersion,
+        venv_path: Option<SystemPathBuf>,
     ) -> Result<Self> {
-        let search_paths = SearchPathSettings::new(src_roots);
+        let mut search_paths = SearchPathSettings::new(src_roots);
+        if let Some(venv_path) = venv_path {
+            search_paths.python_path = PythonPath::from_cli_flag(venv_path);
+        }
 
         let db = Self::default();
         Program::from_settings(