ty: Discover site-packages from the environment that ty is installed in

Author: zanieb

crates/ty/tests/cli/python_environment.rs

@@ -510,6 +510,7 @@ fn many_search_paths() -> anyhow::Result<()> {
     Found 1 diagnostic
 
     ----- stderr -----
+    INFO Using ty's environment for site-packages
     INFO Python version: Python 3.14, platform: linux
     INFO Indexed 7 file(s) in 0.000s
     ");

crates/ty_project/src/metadata/options.rs

@@ -164,14 +164,24 @@ impl Options {
                 .context("Failed to discover local Python environment")?
         };
 
-        let site_packages_paths = if let Some(python_environment) = python_environment.as_ref() {
+        let self_site_packages = self_environment_search_paths(
             python_environment
-                .site_packages_paths(system)
-                .context("Failed to discover the site-packages directory")?
+                .as_ref()
+                .map(ty_python_semantic::PythonEnvironment::origin)
+                .cloned(),
+            system,
+        )
+        .unwrap_or_default();
+
+        let site_packages_paths = if let Some(python_environment) = python_environment.as_ref() {
+            self_site_packages.concatenate(
+                python_environment
+                    .site_packages_paths(system)
+                    .context("Failed to discover the site-packages directory")?,
+            )
         } else {
             tracing::debug!("No virtual environment found");
-
-            SitePackagesPaths::default()
+            self_site_packages
         };
 
         let real_stdlib_path = python_environment.as_ref().and_then(|python_environment| {
@@ -461,6 +471,40 @@ impl Options {
     }
 }
 
+/// Return the site-packages from the environment ty is installed in, as derived from ty's
+/// executable.
+///
+/// If there's an existing environment with an origin that does not allow including site-packages
+/// from ty's environment, discovery of ty's environment is skipped and [`None`] is returned.
+///
+/// Since ty may be executed from an arbitrary non-Python location, errors during discovery of ty's
+/// environment are not raised, instead [`None`] is returned.
+fn self_environment_search_paths(
+    existing_origin: Option<SysPrefixPathOrigin>,
+    system: &dyn System,
+) -> Option<SitePackagesPaths> {
+    if existing_origin.is_some_and(|origin| !origin.allows_extension_with_self_environment()) {
+        return None;
+    }
+
+    let Ok(exe_path) = std::env::current_exe() else {
+        return None;
+    };
+    let ty_path = SystemPath::from_std_path(exe_path.as_path())?;
+
+    let environment = PythonEnvironment::new(ty_path, SysPrefixPathOrigin::SelfEnvironment, system)
+        .inspect_err(|err| tracing::debug!("Failed to discover ty's environment: {err}"))
+        .ok()?;
+
+    tracing::info!("Using ty's environment for site-packages");
+    environment
+        .site_packages_paths(system)
+        .inspect_err(|err| {
+            tracing::debug!("Failed to discover site-packages in ty's environment: {err}");
+        })
+        .ok()
+}
+
 #[derive(
     Debug,
     Default,

crates/ty_python_semantic/src/site_packages.rs

@@ -62,6 +62,15 @@ impl SitePackagesPaths {
         self.0.extend(other.0);
     }
 
+    /// Concatenate two instances of [`SitePackagesPaths`].
+    #[must_use]
+    pub fn concatenate(mut self, other: Self) -> Self {
+        for path in other {
+            self.0.insert(path);
+        }
+        self
+    }
+
     /// Tries to detect the version from the layout of the `site-packages` directory.
     pub fn python_version_from_layout(&self) -> Option<PythonVersionWithSource> {
         if cfg!(windows) {
@@ -252,6 +261,13 @@ impl PythonEnvironment {
             Self::System(env) => env.real_stdlib_directory(system),
         }
     }
+
+    pub fn origin(&self) -> &SysPrefixPathOrigin {
+        match self {
+            Self::Virtual(env) => &env.root_path.origin,
+            Self::System(env) => &env.root_path.origin,
+        }
+    }
 }
 
 /// Enumeration of the subdirectories of `sys.prefix` that could contain a
@@ -1393,15 +1409,15 @@ impl SysPrefixPath {
     ) -> SitePackagesDiscoveryResult<Self> {
         let sys_prefix = if !origin.must_point_directly_to_sys_prefix()
             && system.is_file(unvalidated_path)
-            && unvalidated_path
-                .file_name()
-                .is_some_and(|name| name.starts_with("python"))
-        {
-            // It looks like they passed us a path to a Python executable, e.g. `.venv/bin/python3`.
-            // Try to figure out the `sys.prefix` value from the Python executable.
+            && unvalidated_path.file_name().is_some_and(|name| {
+                name.starts_with("python")
+                    || name.eq_ignore_ascii_case(&format!("ty{}", std::env::consts::EXE_SUFFIX))
+            }) {
+            // It looks like they passed us a path to an executable, e.g. `.venv/bin/python3`. Try
+            // to figure out the `sys.prefix` value from the Python executable.
             let sys_prefix = if cfg!(windows) {
-                // On Windows, the relative path to the Python executable from `sys.prefix`
-                // is different depending on whether it's a virtual environment or a system installation.
+                // On Windows, the relative path to the executable from `sys.prefix` is different
+                // depending on whether it's a virtual environment or a system installation.
                 // System installations have their executable at `<sys.prefix>/python.exe`,
                 // whereas virtual environments have their executable at `<sys.prefix>/Scripts/python.exe`.
                 unvalidated_path.parent().and_then(|parent| {
@@ -1586,6 +1602,8 @@ pub enum SysPrefixPathOrigin {
     /// A `.venv` directory was found in the current working directory,
     /// and the `sys.prefix` path is the path to that virtual environment.
     LocalVenv,
+    /// The `sys.prefix` path came from the environment ty is installed in.
+    SelfEnvironment,
 }
 
 impl SysPrefixPathOrigin {
@@ -1598,6 +1616,7 @@ impl SysPrefixPathOrigin {
             | Self::PythonCliFlag
             | Self::Editor
             | Self::DerivedFromPyvenvCfg
+            | Self::SelfEnvironment
             | Self::CondaPrefixVar => false,
         }
     }
@@ -1608,13 +1627,29 @@ impl SysPrefixPathOrigin {
     /// the `sys.prefix` directory, e.g. the `--python` CLI flag.
     pub(crate) const fn must_point_directly_to_sys_prefix(&self) -> bool {
         match self {
-            Self::PythonCliFlag | Self::ConfigFileSetting(..) | Self::Editor => false,
+            Self::PythonCliFlag
+            | Self::ConfigFileSetting(..)
+            | Self::Editor
+            | Self::SelfEnvironment => false,
             Self::VirtualEnvVar
             | Self::CondaPrefixVar
             | Self::DerivedFromPyvenvCfg
             | Self::LocalVenv => true,
         }
     }
+
+    pub const fn allows_extension_with_self_environment(&self) -> bool {
+        match self {
+            Self::SelfEnvironment
+            | Self::CondaPrefixVar
+            | Self::VirtualEnvVar
+            | Self::Editor
+            | Self::DerivedFromPyvenvCfg
+            | Self::ConfigFileSetting(..)
+            | Self::PythonCliFlag => false,
+            Self::LocalVenv => true,
+        }
+    }
 }
 
 impl std::fmt::Display for SysPrefixPathOrigin {
@@ -1627,6 +1662,7 @@ impl std::fmt::Display for SysPrefixPathOrigin {
             Self::DerivedFromPyvenvCfg => f.write_str("derived `sys.prefix` path"),
             Self::LocalVenv => f.write_str("local virtual environment"),
             Self::Editor => f.write_str("selected interpreter in your editor"),
+            Self::SelfEnvironment => f.write_str("ty environment"),
         }
     }
 }