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

## Summary

This PR does the following things:
- Fixes the `python` configuration setting for mdtest (added in
#17221) so that it expects a path
pointing to a venv's `sys.prefix` variable rather than the a path
pointing to the venv's `site-packages` subdirectory. This brings the
`python` setting in mdtest in sync with our CLI `--python` flag.
- Tweaks mdtest so that it automatically creates a valid `pyvenv.cfg`
file for you if you don't specify one. This makes it much more ergonomic
to write an mdtest with a custom `python` setting: red-knot will reject
a `python` setting that points to a directory that doesn't have a
`pyvenv.cfg` file in it
- Tweaks mdtest so that it doesn't check a custom `pyvenv.cfg` as Python
source code if you _do_ add a custom `pyvenv.cfg` file for your mock
virtual environment in an mdtest. (You get a lot of diagnostics about
Python syntax errors in the `pyvenv.cfg` file, otherwise!)
- Rewrites the test added in
#17178 as an mdtest, and deletes
the original test that was added in that PR

## Test Plan

I verified that the new mdtest fails if I revert the changes to
`resolver.rs` that were added in
#17178

Author: AlexWaygood

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/<path-to-site-packages>/foo/__init__.py`:
+
+```py
+from .a import A as A
+```
+
+`/src/.venv/<path-to-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

@@ -7392,15 +7392,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::*;
@@ -7556,26 +7556,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/README.md

@@ -314,6 +314,45 @@ typeshed = "/typeshed"
 
 For more details, take a look at the [custom-typeshed Markdown test].
 
+### Mocking a virtual environment
+
+Mdtest supports mocking a virtual environment for a specific test at an arbitrary location, again
+using the `[environment]` configuration option:
+
+````markdown
+```toml
+[environment]
+python = ".venv"
+```
+````
+
+Red-knot will reject virtual environments that do not have valid `pyvenv.cfg` files at the
+virtual-environment directory root (here, `.venv/pyvenv.cfg`). However, if a `pyvenv.cfg` file does
+not have its contents specified by the test, mdtest will automatically generate one for you, to
+make mocking a virtual environment more ergonomic.
+
+Mdtest also makes it easy to write Python packages to the mock virtual environment's
+`site-packages` directory using the `<path-to-site-packages>` magic path segment. This would
+otherwise be hard, due to the fact that the `site-packages` subdirectory in a virtual environment
+is located at a different relative path depending on the platform the virtual environment was
+created on. In the following test, mdtest will write the Python file to
+`.venv/Lib/site-packages/foo.py` in its in-memory filesystem used for the test if the test is being
+executed on Windows, and `.venv/lib/python3.13/site-packages/foo.py` otherwise:
+
+````markdown
+```toml
+[environment]
+python = ".venv"
+python-version = "3.13"
+```
+
+`.venv/<path-to-site-packages>/foo.py`:
+
+```py
+X = 1
+```
+````
+
 ## Documentation of tests
 
 Arbitrary Markdown syntax (including of course normal prose paragraphs) is permitted (and ignored by

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,12 @@ fn run_test(
 
     let src_path = project_root.clone();
     let custom_typeshed_path = test.configuration().typeshed();
+    let python_path = test.configuration().python();
+    let python_version = test.configuration().python_version().unwrap_or_default();
+
     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,11 +175,11 @@ 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);
+            let mut full_path = embedded.full_path(&project_root);
 
             if let Some(typeshed_path) = custom_typeshed_path {
                 if let Ok(relative_path) = full_path.strip_prefix(typeshed_path.join("stdlib")) {
@@ -183,11 +189,35 @@ 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;
+                    } else {
+                        let mut new_path = SystemPathBuf::new();
+                        for component in full_path.components() {
+                            let component = component.as_str();
+                            if component == "<path-to-site-packages>" {
+                                if cfg!(target_os = "windows") {
+                                    new_path.push("Lib");
+                                    new_path.push("site-packages");
+                                } else {
+                                    new_path.push("lib");
+                                    new_path.push(format!("python{python_version}"));
+                                    new_path.push("site-packages");
+                                }
+                            } else {
+                                new_path.push(component);
+                            }
+                        }
+                        full_path = new_path;
+                    }
+                }
             }
 
             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,22 +251,34 @@ 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 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 {
-        python_version: configuration.python_version().unwrap_or_default(),
+        python_version,
         python_platform: configuration.python_platform().unwrap_or_default(),
         search_paths: SearchPathSettings {
             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![])),
         },
     };