Update --python to accept paths to executables in virtual environments

Author: zanieb

crates/ty/src/args.rs

@@ -51,14 +51,19 @@ pub(crate) struct CheckCommand {
     #[arg(long, value_name = "PROJECT")]
     pub(crate) project: Option<SystemPathBuf>,
 
-    /// Path to the Python installation from which ty resolves type information and third-party dependencies.
+    /// Path to the Python environment from which ty resolves type information and third-party
+    /// dependencies.
     ///
-    /// If not specified, ty will look at the `VIRTUAL_ENV` environment variable.
+    /// If not specified, ty will attempt to infer it from the `VIRTUAL_ENV` environment variable or
+    /// discover a `.venv` directory in the project root or working directory.
     ///
-    /// ty will search in the path's `site-packages` directories for type information and
-    /// third-party imports.
+    /// If a path to a Python interpreter is provided, e.g., `.venv/bin/python3`, ty will attempt to
+    /// find an environment two directories up from the interpreter's path, e.g., `.venv`. At this
+    /// time, ty does not invoke the interpreter to determine the location of the environment. This
+    /// means that ty will not resolve dynamic executables such as a shim.
     ///
-    /// This option is commonly used to specify the path to a virtual environment.
+    /// ty will search in the given path's `site-packages` directories for type information and
+    /// third-party imports.
     #[arg(long, value_name = "PATH")]
     pub(crate) python: Option<SystemPathBuf>,
 

crates/ty_python_semantic/src/module_resolver/resolver.rs

@@ -247,6 +247,33 @@ impl SearchPaths {
                     .and_then(|env| env.site_packages_directories(system))?
             }
 
+            PythonPath::Resolve(target, origin) => {
+                tracing::debug!("Resolving {origin}: {target}");
+
+                let root = system
+                    // If given a file, assume it's a Python executable, e.g., `.venv/bin/python3`,
+                    // and search for a virtual environment in the root directory. Ideally, we'd
+                    // invoke the target to determine `sys.prefix` here, but that's more complicated
+                    // and may be deferred to uv.
+                    .is_file(target)
+                    .then(|| target.as_path())
+                    .take_if(|target| {
+                        // Avoid using the target if it doesn't look like a Python executable, e.g.,
+                        // to deny cases like `.venv/bin/foo`
+                        target
+                            .file_name()
+                            .is_some_and(|name| name.starts_with("python"))
+                    })
+                    .and_then(SystemPath::parent)
+                    .and_then(SystemPath::parent)
+                    // If not a file, use the path as given and allow let `PythonEnvironment::new`
+                    // handle the error.
+                    .unwrap_or(target);
+
+                PythonEnvironment::new(root, *origin, system)
+                    .and_then(|venv| venv.site_packages_directories(system))?
+            }
+
             PythonPath::Discover(root) => {
                 tracing::debug!("Discovering virtual environment in `{root}`");
                 let virtual_env_path = discover_venv_in(db.system(), root);

crates/ty_python_semantic/src/program.rs

@@ -145,6 +145,9 @@ pub enum PythonPath {
     /// [`sys.prefix`]: https://docs.python.org/3/library/sys.html#sys.prefix
     SysPrefix(SystemPathBuf, SysPrefixPathOrigin),
 
+    /// Resolve a path to an executable (or virtual environment) into a usable environment.
+    Resolve(SystemPathBuf, SysPrefixPathOrigin),
+
     /// Tries to discover a virtual environment in the given path.
     Discover(SystemPathBuf),
 
@@ -161,6 +164,6 @@ impl PythonPath {
     }
 
     pub fn from_cli_flag(path: SystemPathBuf) -> Self {
-        Self::SysPrefix(path, SysPrefixPathOrigin::PythonCliFlag)
+        Self::Resolve(path, SysPrefixPathOrigin::PythonCliFlag)
     }
 }

crates/ty_test/src/config.rs

@@ -68,12 +68,16 @@ pub(crate) struct Environment {
     /// Additional search paths to consider when resolving modules.
     pub(crate) extra_paths: Option<Vec<SystemPathBuf>>,
 
-    /// Path to the Python installation from which ty resolves type information and third-party dependencies.
+    /// Path to the Python installation from which ty resolves type information and third-party
+    /// dependencies.
     ///
-    /// ty will search in the path's `site-packages` directories for type information and
-    /// third-party imports.
+    /// If a path to a Python interpreter is provided, e.g., `.venv/bin/python3`, ty will attempt to
+    /// find an environment two directories up from the interpreter's path, e.g., `.venv`. At this
+    /// time, ty does not invoke the interpreter to determine the location of the environment. This
+    /// means that ty will not resolve dynamic executables such as a shim.
     ///
-    /// This option is commonly used to specify the path to a virtual environment.
+    /// ty will search in the given path's `site-packages` directories for type information and
+    /// third-party imports.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub python: Option<SystemPathBuf>,
 }