[ty] Add PYTHONPATH to EnvVars and fix on Windows (#20490)

Co-authored-by: Micha Reiser <micha@reiser.io>

Cargo.lock

@@ -4291,6 +4291,7 @@ dependencies = [
  "tracing",
  "ty_combine",
  "ty_python_semantic",
+ "ty_static",
  "ty_vendored",
 ]
 

crates/ty/tests/cli/python_environment.rs

@@ -1879,7 +1879,7 @@ fn default_root_python_package_pyi() -> anyhow::Result<()> {
 #[test]
 fn pythonpath_is_respected() -> anyhow::Result<()> {
     let case = CliTest::with_files([
-        ("src/bar/baz.py", "it = 42"),
+        ("baz-dir/baz.py", "it = 42"),
         (
             "src/foo.py",
             r#"
@@ -1915,7 +1915,82 @@ fn pythonpath_is_respected() -> anyhow::Result<()> {
     "#);
 
     assert_cmd_snapshot!(case.command()
-        .env("PYTHONPATH", case.root().join("src/bar")),
+        .env("PYTHONPATH", case.root().join("baz-dir")),
+        @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(())
+}
+
+#[test]
+fn pythonpath_multiple_dirs_is_respected() -> anyhow::Result<()> {
+    let case = CliTest::with_files([
+        ("baz-dir/baz.py", "it = 42"),
+        ("foo-dir/foo.py", "it = 42"),
+        (
+            "src/main.py",
+            r#"
+            import baz
+            import foo
+
+            print(f"{baz.it}")
+            print(f"{foo.it}")
+            "#,
+        ),
+    ])?;
+
+    assert_cmd_snapshot!(case.command(),
+        @r#"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+    error[unresolved-import]: Cannot resolve imported module `baz`
+     --> src/main.py:2:8
+      |
+    2 | import baz
+      |        ^^^
+    3 | import foo
+      |
+    info: Searched in the following paths during module resolution:
+    info:   1. <temp_dir>/ (first-party code)
+    info:   2. <temp_dir>/src (first-party code)
+    info:   3. vendored://stdlib (stdlib typeshed stubs vendored by ty)
+    info: make sure your Python environment is properly configured: https://docs.astral.sh/ty/modules/#python-environment
+    info: rule `unresolved-import` is enabled by default
+
+    error[unresolved-import]: Cannot resolve imported module `foo`
+     --> src/main.py:3:8
+      |
+    2 | import baz
+    3 | import foo
+      |        ^^^
+    4 |
+    5 | print(f"{baz.it}")
+      |
+    info: Searched in the following paths during module resolution:
+    info:   1. <temp_dir>/ (first-party code)
+    info:   2. <temp_dir>/src (first-party code)
+    info:   3. vendored://stdlib (stdlib typeshed stubs vendored by ty)
+    info: make sure your Python environment is properly configured: https://docs.astral.sh/ty/modules/#python-environment
+    info: rule `unresolved-import` is enabled by default
+
+    Found 2 diagnostics
+
+    ----- stderr -----
+    WARN ty is pre-release software and not ready for production use. Expect to encounter bugs, missing features, and fatal errors.
+    "#);
+
+    let pythonpath =
+        std::env::join_paths([case.root().join("baz-dir"), case.root().join("foo-dir")])?;
+    assert_cmd_snapshot!(case.command()
+        .env("PYTHONPATH", pythonpath),
         @r#"
     success: true
     exit_code: 0

crates/ty_project/Cargo.toml

@@ -22,6 +22,7 @@ ruff_python_formatter = { workspace = true, optional = true }
 ruff_text_size = { workspace = true }
 ty_combine = { workspace = true }
 ty_python_semantic = { workspace = true, features = ["serde"] }
+ty_static = { workspace = true }
 ty_vendored = { workspace = true }
 
 anyhow = { workspace = true }

crates/ty_project/src/metadata/options.rs

@@ -34,6 +34,7 @@ use ty_python_semantic::{
     PythonVersionSource, PythonVersionWithSource, SearchPathSettings, SearchPathValidationError,
     SearchPaths, SitePackagesPaths, SysPrefixPathOrigin,
 };
+use ty_static::EnvVars;
 
 #[derive(
     Debug,
@@ -296,38 +297,48 @@ impl Options {
         };
 
         // collect the existing site packages
-        let mut extra_paths: Vec<SystemPathBuf> = Vec::new();
+        let mut extra_paths: Vec<SystemPathBuf> = environment
+            .extra_paths
+            .as_deref()
+            .unwrap_or_default()
+            .iter()
+            .map(|path| path.absolute(project_root, system))
+            .collect();
 
         // read all the paths off the PYTHONPATH environment variable, check
         // they exist as a directory, and add them to the vec of extra_paths
         // as they should be checked before site-packages just like python
         // interpreter does
-        if let Ok(python_path) = system.env_var("PYTHONPATH") {
+        if let Ok(python_path) = system.env_var(EnvVars::PYTHONPATH) {
-            for path in python_path.split(':') {
+            for path in std::env::split_paths(python_path.as_str()) {
-                let possible_path = SystemPath::absolute(path, system.current_directory());
+                let path = match SystemPathBuf::from_path_buf(path) {
+                    Ok(path) => path,
+                    Err(path) => {
+                        tracing::debug!(
+                            "Skipping `{path}` listed in `PYTHONPATH` because the path is not valid UTF-8",
+                            path = path.display()
+                        );
+                        continue;
+                    }
+                };
 
-                if system.is_directory(&possible_path) {
+                let abspath = SystemPath::absolute(path, system.current_directory());
+
+                if !system.is_directory(&abspath) {
                     tracing::debug!(
-                        "Adding `{possible_path}` from the `PYTHONPATH` environment variable to `extra_paths`"
+                        "Skipping `{abspath}` listed in `PYTHONPATH` because the path doesn't exist or isn't a directory"
-                    );
-                    extra_paths.push(possible_path);
-                } else {
-                    tracing::debug!(
-                        "Skipping `{possible_path}` listed in `PYTHONPATH` because the path doesn't exist or isn't a directory"
                     );
+                    continue;
                 }
+
+                tracing::debug!(
+                    "Adding `{abspath}` from the `PYTHONPATH` environment variable to `extra_paths`"
+                );
+
+                extra_paths.push(abspath);
             }
         }
 
-        extra_paths.extend(
-            environment
-                .extra_paths
-                .as_deref()
-                .unwrap_or_default()
-                .iter()
-                .map(|path| path.absolute(project_root, system)),
-        );
-
         let settings = SearchPathSettings {
             extra_paths,
             src_roots,

crates/ty_static/src/env_vars.rs

@@ -42,6 +42,12 @@ impl EnvVars {
     /// Used to detect an activated virtual environment.
     pub const VIRTUAL_ENV: &'static str = "VIRTUAL_ENV";
 
+    /// Adds additional directories to ty's search paths.
+    /// The format is the same as the shell’s PATH:
+    /// one or more directory pathnames separated by os appropriate pathsep
+    /// (e.g. colons on Unix or semicolons on Windows).
+    pub const PYTHONPATH: &'static str = "PYTHONPATH";
+
     /// Used to determine if an active Conda environment is the base environment or not.
     pub const CONDA_DEFAULT_ENV: &'static str = "CONDA_DEFAULT_ENV";