[ty] Resolving Python path using CONDA_PREFIX variable to support Conda and Pixi (#18267)

Author: lipefree

crates/ty/tests/cli.rs

@@ -1555,6 +1555,83 @@ fn cli_config_args_invalid_option() -> anyhow::Result<()> {
     Ok(())
 }
 
+/// The `site-packages` directory is used by ty for external import.
+/// Ty does the following checks to discover the `site-packages` directory in the order:
+/// 1) If `VIRTUAL_ENV` environment variable is set
+/// 2) If `CONDA_PREFIX` environment variable is set
+/// 3) If a `.venv` directory exists at the project root
+///
+/// This test is aiming at validating the logic around `CONDA_PREFIX`.
+///
+/// A conda-like environment file structure is used
+/// We test by first not setting the `CONDA_PREFIX` and expect a fail.
+/// Then we test by setting `CONDA_PREFIX` to `conda-env` and expect a pass.
+///
+/// ├── project
+/// │   └── test.py
+/// └── conda-env
+///     └── lib
+///         └── python3.13
+///             └── site-packages
+///                 └── package1
+///                     └── __init__.py
+///
+/// test.py imports package1
+/// And the command is run in the `project` directory.
+#[test]
+fn check_conda_prefix_var_to_resolve_path() -> anyhow::Result<()> {
+    let conda_package1_path = if cfg!(windows) {
+        "conda-env/Lib/site-packages/package1/__init__.py"
+    } else {
+        "conda-env/lib/python3.13/site-packages/package1/__init__.py"
+    };
+
+    let case = TestCase::with_files([
+        (
+            "project/test.py",
+            r#"
+            import package1
+            "#,
+        ),
+        (
+            conda_package1_path,
+            r#"
+            "#,
+        ),
+    ])?;
+
+    assert_cmd_snapshot!(case.command().current_dir(case.root().join("project")), @r"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+    error[unresolved-import]: Cannot resolve imported module `package1`
+     --> test.py:2:8
+      |
+    2 | import package1
+      |        ^^^^^^^^
+      |
+    info: make sure your Python environment is properly configured: https://github.com/astral-sh/ty/blob/main/docs/README.md#python-environment
+    info: rule `unresolved-import` is enabled by default
+
+    Found 1 diagnostic
+
+    ----- stderr -----
+    WARN ty is pre-release software and not ready for production use. Expect to encounter bugs, missing features, and fatal errors.
+    ");
+
+    // do command : CONDA_PREFIX=<temp_dir>/conda_env
+    assert_cmd_snapshot!(case.command().current_dir(case.root().join("project")).env("CONDA_PREFIX", case.root().join("conda-env")), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    All checks passed!
+
+    ----- stderr -----
+    WARN ty is pre-release software and not ready for production use. Expect to encounter bugs, missing features, and fatal errors.
+    ");
+    Ok(())
+}
+
 struct TestCase {
     _temp_dir: TempDir,
     _settings_scope: SettingsBindDropGuard,

crates/ty_project/src/metadata/options.rs

@@ -186,6 +186,11 @@ impl Options {
                         .ok()
                         .map(PythonPath::from_virtual_env_var)
                 })
+                .or_else(|| {
+                    std::env::var("CONDA_PREFIX")
+                        .ok()
+                        .map(PythonPath::from_conda_prefix_var)
+                })
                 .unwrap_or_else(|| PythonPath::Discover(project_root.to_path_buf())),
         }
     }

crates/ty_python_semantic/src/program.rs

@@ -202,6 +202,10 @@ impl PythonPath {
         Self::SysPrefix(path.into(), SysPrefixPathOrigin::VirtualEnvVar)
     }
 
+    pub fn from_conda_prefix_var(path: impl Into<SystemPathBuf>) -> Self {
+        Self::Resolve(path.into(), SysPrefixPathOrigin::CondaPrefixVar)
+    }
+
     pub fn from_cli_flag(path: SystemPathBuf) -> Self {
         Self::Resolve(path, SysPrefixPathOrigin::PythonCliFlag)
     }

crates/ty_python_semantic/src/site_packages.rs

@@ -580,6 +580,7 @@ impl fmt::Display for SysPrefixPath {
 pub enum SysPrefixPathOrigin {
     PythonCliFlag,
     VirtualEnvVar,
+    CondaPrefixVar,
     Derived,
     LocalVenv,
 }
@@ -590,7 +591,7 @@ impl SysPrefixPathOrigin {
     pub(crate) fn must_be_virtual_env(self) -> bool {
         match self {
             Self::LocalVenv | Self::VirtualEnvVar => true,
-            Self::PythonCliFlag | Self::Derived => false,
+            Self::PythonCliFlag | Self::Derived | Self::CondaPrefixVar => false,
         }
     }
 }
@@ -600,6 +601,7 @@ impl Display for SysPrefixPathOrigin {
         match self {
             Self::PythonCliFlag => f.write_str("`--python` argument"),
             Self::VirtualEnvVar => f.write_str("`VIRTUAL_ENV` environment variable"),
+            Self::CondaPrefixVar => f.write_str("`CONDA_PREFIX` environment variable"),
             Self::Derived => f.write_str("derived `sys.prefix` path"),
             Self::LocalVenv => f.write_str("local virtual environment"),
         }