[ty] Add option for ty to find modules in paths defined in PYTHONPATH

crates/ty_project/src/metadata/options.rs

@@ -22,6 +22,7 @@ use ruff_python_ast::PythonVersion;
 use rustc_hash::FxHasher;
 use serde::{Deserialize, Serialize};
 use std::borrow::Cow;
+use std::env;
 use std::fmt::{self, Debug, Display};
 use std::hash::BuildHasherDefault;
 use std::ops::Deref;
@@ -277,6 +278,28 @@ impl Options {
             roots
         };
 
+        // collect the existing site packages
+        let mut site_packages_paths = site_packages_paths.into_vec();
+
+        // If read_python_path is defined in config, read all the paths off the
+        // PYTHONPATH environment variable, check they exists, and add them to
+        // the vec of site_package_paths, as that is conceptually what they are.
+        if environment.read_python_path.unwrap_or_default() {
+            if let Ok(python_path) = env::var("PYTHONPATH") {
+                let splits: Vec<SystemPathBuf> = python_path
+                    .split(':')
+                    .filter_map(|str_path| {
+                        let possible_path = std::path::PathBuf::from(str_path);
+                        match possible_path.try_exists() {
+                            Ok(true) => SystemPathBuf::from_path_buf(possible_path).ok(),
+                            _ => None,
+                        }
+                    })
+                    .collect();
+                site_packages_paths.extend(splits);
+            }
+        }
+
         let settings = SearchPathSettings {
             extra_paths: environment
                 .extra_paths
@@ -290,7 +313,7 @@ impl Options {
                 .typeshed
                 .as_ref()
                 .map(|path| path.absolute(project_root, system)),
-            site_packages_paths: site_packages_paths.into_vec(),
+            site_packages_paths,
             real_stdlib_path,
         };
 
@@ -441,6 +464,17 @@ pub struct EnvironmentOptions {
     )]
     pub root: Option<Vec<RelativePathBuf>>,
 
+    /// Read PYTHONPATH to find modules
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[option(
+        default = r#"false"#,
+        value_type = "bool",
+        example = r#"
+            read_python_path = true
+        "#
+    )]
+    pub read_python_path: Option<bool>,
+
     /// Specifies the version of Python that will be used to analyze the source code.
     /// The version should be specified as a string in the format `M.m` where `M` is the major version
     /// and `m` is the minor (e.g. `"3.0"` or `"3.6"`).

crates/ty_python_semantic/src/module_resolver/resolver.rs

@@ -675,6 +675,8 @@ fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Opti
     let stub_name = name.to_stub_package();
     let mut is_namespace_package = false;
 
+    // Make an stub tracker variable. Don't return None until all paths have been checked.
+    let mut found_stub_without_module = false;
     for search_path in search_paths(db, mode) {
         // When a builtin module is imported, standard module resolution is bypassed:
         // the module name always resolves to the stdlib module,
@@ -718,7 +720,10 @@ fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Opti
                         "Stub-package in `{search_path}` doesn't contain module: `{name}`"
                     );
                     // stub exists, but the module doesn't.
-                    return None;
+                    // mark the failure, but don't fail yet, wait for all paths
+                    // to be checked.
+                    found_stub_without_module = true;
+                    continue;
                 }
                 Err((PackageKind::Namespace, _)) => {
                     tracing::trace!(
@@ -754,7 +759,11 @@ fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Opti
                     // For regular packages, don't search the next search path. All files of that
                     // package must be in the same location
                     tracing::trace!("Package in `{search_path}` doesn't contain module: `{name}`");
-                    return None;
+                    // stub exists, but the module doesn't.
+                    // mark the failure, but don't fail yet, wait for all paths
+                    // to be checked.
+                    found_stub_without_module = true;
+                    continue;
                 }
                 (PackageKind::Namespace, _) => {
                     tracing::trace!(
@@ -765,6 +774,12 @@ fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Opti
             },
         }
     }
+    if found_stub_without_module {
+        // There were one or more paths that had part of the module stub but the module
+        // was not found in any of them.
+        tracing::trace!("No search path contains module: `{name}`");
+        return None;
+    }
 
     if is_namespace_package {
         return Some(ResolvedName::NamespacePackage);