[ty] Search PYTHONPATH to find modules (#20441)

Co-authored-by: Nate Lust <[email protected]>
Co-authored-by: Micha Reiser <[email protected]>

Author: 3 people

crates/ty/tests/cli/python_environment.rs

@@ -1875,3 +1875,56 @@ fn default_root_python_package_pyi() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn pythonpath_is_respected() -> anyhow::Result<()> {
+    let case = CliTest::with_files([
+        ("src/bar/baz.py", "it = 42"),
+        (
+            "src/foo.py",
+            r#"
+            import baz
+            print(f"{baz.it}")
+            "#,
+        ),
+    ])?;
+
+    assert_cmd_snapshot!(case.command(),
+        @r#"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+    error[unresolved-import]: Cannot resolve imported module `baz`
+     --> src/foo.py:2:8
+      |
+    2 | import baz
+      |        ^^^
+    3 | 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 1 diagnostic
+
+    ----- stderr -----
+    WARN ty is pre-release software and not ready for production use. Expect to encounter bugs, missing features, and fatal errors.
+    "#);
+
+    assert_cmd_snapshot!(case.command()
+        .env("PYTHONPATH", case.root().join("src/bar")),
+        @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(())
+}

crates/ty_project/src/metadata/options.rs

@@ -295,14 +295,41 @@ impl Options {
             roots
         };
 
-        let settings = SearchPathSettings {
-            extra_paths: environment
+        // collect the existing site packages
+        let mut extra_paths: Vec<SystemPathBuf> = Vec::new();
+
+        // 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") {
+            for path in python_path.split(':') {
+                let possible_path = SystemPath::absolute(path, system.current_directory());
+
+                if system.is_directory(&possible_path) {
+                    tracing::debug!(
+                        "Adding `{possible_path}` from the `PYTHONPATH` environment variable to `extra_paths`"
+                    );
+                    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"
+                    );
+                }
+            }
+        }
+
+        extra_paths.extend(
+            environment
                 .extra_paths
                 .as_deref()
                 .unwrap_or_default()
                 .iter()
-                .map(|path| path.absolute(project_root, system))
-                .collect(),
+                .map(|path| path.absolute(project_root, system)),
+        );
+
+        let settings = SearchPathSettings {
+            extra_paths,
             src_roots,
             custom_typeshed: environment
                 .typeshed