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

crates/ty/tests/cli/main.rs

@@ -5,6 +5,7 @@ mod python_environment;
 mod rule_selection;
 
 use anyhow::Context as _;
+use insta::Settings;
 use insta::internals::SettingsBindDropGuard;
 use insta_cmd::{assert_cmd_snapshot, get_cargo_bin};
 use std::{
@@ -760,8 +761,10 @@ fn can_handle_large_binop_expressions() -> anyhow::Result<()> {
 
 pub(crate) struct CliTest {
     _temp_dir: TempDir,
-    _settings_scope: SettingsBindDropGuard,
+    settings: Settings,
+    settings_scope: Option<SettingsBindDropGuard>,
     project_dir: PathBuf,
+    ty_binary_path: PathBuf,
 }
 
 impl CliTest {
@@ -794,7 +797,9 @@ impl CliTest {
         Ok(Self {
             project_dir,
             _temp_dir: temp_dir,
-            _settings_scope: settings_scope,
+            settings,
+            settings_scope: Some(settings_scope),
+            ty_binary_path: get_cargo_bin("ty"),
         })
     }
 
@@ -823,6 +828,30 @@ impl CliTest {
         Ok(())
     }
 
+    /// Return [`Self`] with the ty binary copied to the specified path instead.
+    pub(crate) fn with_ty_at(mut self, dest_path: impl AsRef<Path>) -> anyhow::Result<Self> {
+        let dest_path = dest_path.as_ref();
+        let dest_path = self.project_dir.join(dest_path);
+
+        Self::ensure_parent_directory(&dest_path)?;
+        std::fs::copy(&self.ty_binary_path, &dest_path)
+            .with_context(|| format!("Failed to copy ty binary to `{}`", dest_path.display()))?;
+
+        self.ty_binary_path = dest_path;
+        Ok(self)
+    }
+
+    /// Add a filter to the settings and rebind them.
+    pub(crate) fn with_filter(mut self, pattern: &str, replacement: &str) -> Self {
+        self.settings.add_filter(pattern, replacement);
+        // Drop the old scope before binding a new one, otherwise the old scope is dropped _after_
+        // binding and assigning the new one, restoring the settings to their state before the old
+        // scope was bound.
+        drop(self.settings_scope.take());
+        self.settings_scope = Some(self.settings.bind_to_scope());
+        self
+    }
+
     fn ensure_parent_directory(path: &Path) -> anyhow::Result<()> {
         if let Some(parent) = path.parent() {
             std::fs::create_dir_all(parent)
@@ -868,7 +897,7 @@ impl CliTest {
     }
 
     pub(crate) fn command(&self) -> Command {
-        let mut command = Command::new(get_cargo_bin("ty"));
+        let mut command = Command::new(&self.ty_binary_path);
         command.current_dir(&self.project_dir).arg("check");
 
         // Unset all environment variables because they can affect test behavior.
@@ -881,3 +910,11 @@ impl CliTest {
 fn tempdir_filter(path: &Path) -> String {
     format!(r"{}\\?/?", regex::escape(path.to_str().unwrap()))
 }
+
+fn site_packages_filter(python_version: &str) -> String {
+    if cfg!(windows) {
+        "Lib/site-packages".to_string()
+    } else {
+        format!("lib/python{}/site-packages", regex::escape(python_version))
+    }
+}

crates/ty/tests/cli/python_environment.rs

@@ -1,7 +1,7 @@
 use insta_cmd::assert_cmd_snapshot;
 use ruff_python_ast::PythonVersion;
 
-use crate::CliTest;
+use crate::{CliTest, site_packages_filter};
 
 /// Specifying an option on the CLI should take precedence over the same setting in the
 /// project's configuration. Here, this is tested for the Python version.
@@ -1654,6 +1654,278 @@ home = ./
     Ok(())
 }
 
+/// ty should include site packages from its own environment when no other environment is found.
+#[test]
+fn ty_environment_is_only_environment() -> anyhow::Result<()> {
+    let ty_venv_site_packages = if cfg!(windows) {
+        "ty-venv/Lib/site-packages"
+    } else {
+        "ty-venv/lib/python3.13/site-packages"
+    };
+
+    let ty_executable_path = if cfg!(windows) {
+        "ty-venv/Scripts/ty.exe"
+    } else {
+        "ty-venv/bin/ty"
+    };
+
+    let ty_package_path = format!("{ty_venv_site_packages}/ty_package/__init__.py");
+
+    let case = CliTest::with_files([
+        (ty_package_path.as_str(), "class TyEnvClass: ..."),
+        (
+            "ty-venv/pyvenv.cfg",
+            r"
+            home = ./
+            version = 3.13
+            ",
+        ),
+        (
+            "test.py",
+            r"
+            from ty_package import TyEnvClass
+            ",
+        ),
+    ])?;
+
+    let case = case.with_ty_at(ty_executable_path)?;
+    assert_cmd_snapshot!(case.command(), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    All checks passed!
+
+    ----- stderr -----
+    "###);
+
+    Ok(())
+}
+
+/// ty should include site packages from both its own environment and a local `.venv`. The packages
+/// from ty's environment should take precedence.
+#[test]
+fn ty_environment_and_discovered_venv() -> anyhow::Result<()> {
+    let ty_venv_site_packages = if cfg!(windows) {
+        "ty-venv/Lib/site-packages"
+    } else {
+        "ty-venv/lib/python3.13/site-packages"
+    };
+
+    let ty_executable_path = if cfg!(windows) {
+        "ty-venv/Scripts/ty.exe"
+    } else {
+        "ty-venv/bin/ty"
+    };
+
+    let local_venv_site_packages = if cfg!(windows) {
+        ".venv/Lib/site-packages"
+    } else {
+        ".venv/lib/python3.13/site-packages"
+    };
+
+    let ty_unique_package = format!("{ty_venv_site_packages}/ty_package/__init__.py");
+    let local_unique_package = format!("{local_venv_site_packages}/local_package/__init__.py");
+    let ty_conflicting_package = format!("{ty_venv_site_packages}/shared_package/__init__.py");
+    let local_conflicting_package =
+        format!("{local_venv_site_packages}/shared_package/__init__.py");
+
+    let case = CliTest::with_files([
+        (ty_unique_package.as_str(), "class TyEnvClass: ..."),
+        (local_unique_package.as_str(), "class LocalClass: ..."),
+        (ty_conflicting_package.as_str(), "class FromTyEnv: ..."),
+        (
+            local_conflicting_package.as_str(),
+            "class FromLocalVenv: ...",
+        ),
+        (
+            "ty-venv/pyvenv.cfg",
+            r"
+            home = ./
+            version = 3.13
+            ",
+        ),
+        (
+            ".venv/pyvenv.cfg",
+            r"
+            home = ./
+            version = 3.13
+            ",
+        ),
+        (
+            "test.py",
+            r"
+            # Should resolve from ty's environment
+            from ty_package import TyEnvClass
+            # Should resolve from local .venv
+            from local_package import LocalClass
+            # Should resolve from ty's environment (takes precedence)
+            from shared_package import FromTyEnv
+            # Should NOT resolve (shadowed by ty's environment version)
+            from shared_package import FromLocalVenv
+            ",
+        ),
+    ])?
+    .with_ty_at(ty_executable_path)?;
+
+    assert_cmd_snapshot!(case.command(), @r###"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+    error[unresolved-import]: Module `shared_package` has no member `FromLocalVenv`
+     --> test.py:9:28
+      |
+    7 | from shared_package import FromTyEnv
+    8 | # Should NOT resolve (shadowed by ty's environment version)
+    9 | from shared_package import FromLocalVenv
+      |                            ^^^^^^^^^^^^^
+      |
+    info: rule `unresolved-import` is enabled by default
+
+    Found 1 diagnostic
+
+    ----- stderr -----
+    "###);
+
+    Ok(())
+}
+
+/// When `VIRTUAL_ENV` is set, ty should *not* discover its own environment's site-packages.
+#[test]
+fn ty_environment_and_active_environment() -> anyhow::Result<()> {
+    let ty_venv_site_packages = if cfg!(windows) {
+        "ty-venv/Lib/site-packages"
+    } else {
+        "ty-venv/lib/python3.13/site-packages"
+    };
+
+    let ty_executable_path = if cfg!(windows) {
+        "ty-venv/Scripts/ty.exe"
+    } else {
+        "ty-venv/bin/ty"
+    };
+
+    let active_venv_site_packages = if cfg!(windows) {
+        "active-venv/Lib/site-packages"
+    } else {
+        "active-venv/lib/python3.13/site-packages"
+    };
+
+    let ty_package_path = format!("{ty_venv_site_packages}/ty_package/__init__.py");
+    let active_package_path = format!("{active_venv_site_packages}/active_package/__init__.py");
+
+    let case = CliTest::with_files([
+        (ty_package_path.as_str(), "class TyEnvClass: ..."),
+        (
+            "ty-venv/pyvenv.cfg",
+            r"
+            home = ./
+            version = 3.13
+            ",
+        ),
+        (active_package_path.as_str(), "class ActiveClass: ..."),
+        (
+            "active-venv/pyvenv.cfg",
+            r"
+            home = ./
+            version = 3.13
+            ",
+        ),
+        (
+            "test.py",
+            r"
+            from ty_package import TyEnvClass
+            from active_package import ActiveClass
+            ",
+        ),
+    ])?
+    .with_ty_at(ty_executable_path)?
+    .with_filter(&site_packages_filter("3.13"), "<site-packages>");
+
+    assert_cmd_snapshot!(
+        case.command()
+            .env("VIRTUAL_ENV", case.root().join("active-venv")),
+        @r"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+    error[unresolved-import]: Cannot resolve imported module `ty_package`
+     --> test.py:2:6
+      |
+    2 | from ty_package import TyEnvClass
+      |      ^^^^^^^^^^
+    3 | from active_package import ActiveClass
+      |
+    info: Searched in the following paths during module resolution:
+    info:   1. <temp_dir>/ (first-party code)
+    info:   2. vendored://stdlib (stdlib typeshed stubs vendored by ty)
+    info:   3. <temp_dir>/active-venv/<site-packages> (site-packages)
+    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 -----
+    "
+    );
+
+    Ok(())
+}
+
+/// When ty is installed in a system environment rather than a virtual environment, it should
+/// not include the environment's site-packages in its search path.
+#[test]
+fn ty_environment_is_system_not_virtual() -> anyhow::Result<()> {
+    let ty_system_site_packages = if cfg!(windows) {
+        "system-python/Lib/site-packages"
+    } else {
+        "system-python/lib/python3.13/site-packages"
+    };
+
+    let ty_executable_path = if cfg!(windows) {
+        "system-python/Scripts/ty.exe"
+    } else {
+        "system-python/bin/ty"
+    };
+
+    let ty_package_path = format!("{ty_system_site_packages}/system_package/__init__.py");
+
+    let case = CliTest::with_files([
+        // Package in system Python installation (should NOT be discovered)
+        (ty_package_path.as_str(), "class SystemClass: ..."),
+        // Note: NO pyvenv.cfg - this is a system installation, not a venv
+        (
+            "test.py",
+            r"
+            from system_package import SystemClass
+            ",
+        ),
+    ])?
+    .with_ty_at(ty_executable_path)?;
+
+    assert_cmd_snapshot!(case.command(), @r###"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+    error[unresolved-import]: Cannot resolve imported module `system_package`
+     --> test.py:2:6
+      |
+    2 | from system_package import SystemClass
+      |      ^^^^^^^^^^^^^^
+      |
+    info: Searched in the following paths during module resolution:
+    info:   1. <temp_dir>/ (first-party code)
+    info:   2. 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 -----
+    "###);
+
+    Ok(())
+}
+
 #[test]
 fn src_root_deprecation_warning() -> anyhow::Result<()> {
     let case = CliTest::with_files([