[red-knot] Fix python setting in mdtests, and rewrite a site-packages test as an mdtest

crates/red_knot_python_semantic/resources/mdtest/import/relative.md

@@ -236,3 +236,36 @@ X: int = 42
 ```py
 from .parser import X  # error: [unresolved-import]
 ```
+
+## Relative imports in `site-packages`
+
+Relative imports in `site-packages` are correctly resolved even when the `site-packages` search path
+is a subdirectory of the first-party search path. Note that mdtest sets the first-party search path
+to `/src/`, which is why the virtual environment in this test is a subdirectory of `/src/`, even
+though this is not how a typical Python project would be structured:
+
+```toml
+[environment]
+python = "/src/.venv"
+python-version = "3.13"
+```
+
+`/src/bar.py`:
+
+```py
+from foo import A
+
+reveal_type(A)  # revealed: Literal[A]
+```
+
+`/src/.venv/lib/python3.13/site-packages/foo/__init__.py`:
+
+```py
+from .a import A as A
+```
+
+`/src/.venv/lib/python3.13/site-packages/foo/a.py`:
+
+```py
+class A: ...
+```

crates/red_knot_python_semantic/src/db.rs

@@ -19,7 +19,7 @@ pub(crate) mod tests {
     use std::sync::Arc;
 
     use crate::program::{Program, SearchPathSettings};
-    use crate::{default_lint_registry, ProgramSettings, PythonPath, PythonPlatform};
+    use crate::{default_lint_registry, ProgramSettings, PythonPlatform};
 
     use super::Db;
     use crate::lint::{LintRegistry, RuleSelection};
@@ -139,8 +139,6 @@ pub(crate) mod tests {
         python_version: PythonVersion,
         /// Target Python platform
         python_platform: PythonPlatform,
-        /// Paths to the directory to use for `site-packages`
-        site_packages: Vec<SystemPathBuf>,
         /// Path and content pairs for files that should be present
         files: Vec<(&'a str, &'a str)>,
     }
@@ -150,7 +148,6 @@ pub(crate) mod tests {
             Self {
                 python_version: PythonVersion::default(),
                 python_platform: PythonPlatform::default(),
-                site_packages: vec![],
                 files: vec![],
             }
         }
@@ -169,14 +166,6 @@ pub(crate) mod tests {
             self
         }
 
-        pub(crate) fn with_site_packages_search_path(
-            mut self,
-            path: &(impl AsRef<SystemPath> + ?Sized),
-        ) -> Self {
-            self.site_packages.push(path.as_ref().to_path_buf());
-            self
-        }
-
         pub(crate) fn build(self) -> anyhow::Result<TestDb> {
             let mut db = TestDb::new();
 
@@ -186,15 +175,12 @@ pub(crate) mod tests {
             db.write_files(self.files)
                 .context("Failed to write test files")?;
 
-            let mut search_paths = SearchPathSettings::new(vec![src_root]);
-            search_paths.python_path = PythonPath::KnownSitePackages(self.site_packages);
-
             Program::from_settings(
                 &db,
                 ProgramSettings {
                     python_version: self.python_version,
                     python_platform: self.python_platform,
-                    search_paths,
+                    search_paths: SearchPathSettings::new(vec![src_root]),
                 },
             )
             .context("Failed to configure Program settings")?;

crates/red_knot_python_semantic/src/lib.rs

@@ -10,6 +10,7 @@ pub use module_resolver::{resolve_module, system_module_search_paths, KnownModul
 pub use program::{Program, ProgramSettings, PythonPath, SearchPathSettings};
 pub use python_platform::PythonPlatform;
 pub use semantic_model::{HasType, SemanticModel};
+pub use site_packages::SysPrefixPathOrigin;
 
 pub mod ast_node_ref;
 mod db;

crates/red_knot_python_semantic/src/types/infer.rs

@@ -7386,15 +7386,15 @@ impl StringPartsCollector {
 
 #[cfg(test)]
 mod tests {
-    use crate::db::tests::{setup_db, TestDb, TestDbBuilder};
+    use crate::db::tests::{setup_db, TestDb};
     use crate::semantic_index::definition::Definition;
     use crate::semantic_index::symbol::FileScopeId;
     use crate::semantic_index::{global_scope, semantic_index, symbol_table, use_def_map};
     use crate::symbol::global_symbol;
     use crate::types::check_types;
     use ruff_db::diagnostic::Diagnostic;
     use ruff_db::files::{system_path_to_file, File};
-    use ruff_db::system::{DbWithWritableSystem as _, SystemPath};
+    use ruff_db::system::DbWithWritableSystem as _;
     use ruff_db::testing::{assert_function_query_was_not_run, assert_function_query_was_run};
 
     use super::*;
@@ -7550,26 +7550,6 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn relative_import_resolution_in_site_packages_when_site_packages_is_subdirectory_of_first_party_search_path(
-    ) {
-        let project_root = SystemPath::new("/src");
-        let foo_dot_py = project_root.join("foo.py");
-        let site_packages = project_root.join(".venv/lib/python3.13/site-packages");
-
-        let db = TestDbBuilder::new()
-            .with_site_packages_search_path(&site_packages)
-            .with_file(&foo_dot_py, "from bar import A")
-            .with_file(&site_packages.join("bar/__init__.py"), "from .a import *")
-            .with_file(&site_packages.join("bar/a.py"), "class A: ...")
-            .build()
-            .unwrap();
-
-        assert_file_diagnostics(&db, foo_dot_py.as_str(), &[]);
-        let a_symbol = get_symbol(&db, foo_dot_py.as_str(), &[], "A");
-        assert!(a_symbol.expect_type().is_class_literal());
-    }
-
     #[test]
     fn pep695_type_params() {
         let mut db = setup_db();

crates/red_knot_test/src/lib.rs

@@ -5,7 +5,9 @@ use colored::Colorize;
 use config::SystemKind;
 use parser as test_parser;
 use red_knot_python_semantic::types::check_types;
-use red_knot_python_semantic::{Program, ProgramSettings, PythonPath, SearchPathSettings};
+use red_knot_python_semantic::{
+    Program, ProgramSettings, PythonPath, SearchPathSettings, SysPrefixPathOrigin,
+};
 use ruff_db::diagnostic::{create_parse_diagnostic, Diagnostic, DisplayDiagnosticConfig};
 use ruff_db::files::{system_path_to_file, File};
 use ruff_db::panic::catch_unwind;
@@ -158,8 +160,10 @@ fn run_test(
 
     let src_path = project_root.clone();
     let custom_typeshed_path = test.configuration().typeshed();
+    let python_path = test.configuration().python();
     let mut typeshed_files = vec![];
     let mut has_custom_versions_file = false;
+    let mut has_custom_pyvenv_cfg_file = false;
 
     let test_files: Vec<_> = test
         .files()
@@ -169,8 +173,8 @@ fn run_test(
             }
 
             assert!(
-                matches!(embedded.lang, "py" | "pyi" | "python" | "text"),
-                "Supported file types are: py (or python), pyi, text, and ignore"
+                matches!(embedded.lang, "py" | "pyi" | "python" | "text" | "cfg"),
+                "Supported file types are: py (or python), pyi, text, cfg and ignore"
             );
 
             let full_path = embedded.full_path(&project_root);
@@ -183,11 +187,17 @@ fn run_test(
                         typeshed_files.push(relative_path.to_path_buf());
                     }
                 }
+            } else if let Some(python_path) = python_path {
+                if let Ok(relative_path) = full_path.strip_prefix(python_path) {
+                    if relative_path.as_str() == "pyvenv.cfg" {
+                        has_custom_pyvenv_cfg_file = true;
+                    }
+                }
             }
 
             db.write_file(&full_path, &embedded.code).unwrap();
 
-            if !full_path.starts_with(&src_path) || embedded.lang == "text" {
+            if !full_path.starts_with(&src_path) || !matches!(embedded.lang, "py" | "pyi") {
                 // These files need to be written to the file system (above), but we don't run any checks on them.
                 return None;
             }
@@ -221,6 +231,17 @@ fn run_test(
         }
     }
 
+    if let Some(python_path) = python_path {
+        if !has_custom_pyvenv_cfg_file {
+            let pyvenv_cfg_file = python_path.join("pyvenv.cfg");
+            let python_version = test.configuration().python_version().unwrap_or_default();
+            let home_directory = SystemPathBuf::from(format!("/Python{python_version}"));
+            db.create_directory_all(&home_directory).unwrap();
+            db.write_file(&pyvenv_cfg_file, format!("home = {home_directory}"))
+                .unwrap();
+        }
+    }
+
     let configuration = test.configuration();
 
     let settings = ProgramSettings {
@@ -230,13 +251,15 @@ fn run_test(
             src_roots: vec![src_path],
             extra_paths: configuration.extra_paths().unwrap_or_default().to_vec(),
             custom_typeshed: custom_typeshed_path.map(SystemPath::to_path_buf),
-            python_path: PythonPath::KnownSitePackages(
-                configuration
-                    .python()
-                    .into_iter()
-                    .map(SystemPath::to_path_buf)
-                    .collect(),
-            ),
+            python_path: configuration
+                .python()
+                .map(|sys_prefix| {
+                    PythonPath::SysPrefix(
+                        sys_prefix.to_path_buf(),
+                        SysPrefixPathOrigin::PythonCliFlag,
+                    )
+                })
+                .unwrap_or(PythonPath::KnownSitePackages(vec![])),
         },
     };