1use ruff_db::files::File;
2use ty_module_resolver::Db as ModuleResolverDb;
3
4#[cfg(any(test, feature = "testing"))]
5use crate::program::{Program, ProgramSettings};
6
7#[salsa::db]
9pub trait Db: ModuleResolverDb {
10 fn should_check_file(&self, file: File) -> bool;
12}
13
14#[cfg(any(test, feature = "testing"))]
15#[salsa::db]
16pub trait TestProgramDb: Db {
17 fn program_settings(&self) -> &ProgramSettings;
18
19 fn program(&self) -> Program<'_>
21 where
22 Self: Sized,
23 {
24 #[salsa::tracked(returns(copy), heap_size=ruff_memory_usage::heap_size)]
25 fn program_inner(db: &dyn TestProgramDb) -> Program<'_> {
26 Program::from_settings(db, db.program_settings())
27 }
28
29 program_inner(self)
30 }
31}
32
33#[cfg(test)]
34pub(crate) mod tests {
35 use std::sync::{Arc, Mutex};
36
37 use anyhow::Context;
38
39 use ruff_db::Db as SourceDb;
40 use ruff_db::files::{File, Files};
41 use ruff_db::system::{
42 DbWithTestSystem, DbWithWritableSystem as _, System, SystemPath, SystemPathBuf, TestSystem,
43 };
44 use ruff_db::vendored::VendoredFileSystem;
45 use ruff_python_ast::PythonVersion;
46 use ty_module_resolver::{Db as ModuleResolverDb, FallibleStrategy, SearchPathSettings};
47 use ty_site_packages::{PythonVersionSource, PythonVersionWithSource};
48
49 use crate::platform::PythonPlatform;
50 use crate::program::ProgramSettings;
51
52 use super::{Db, TestProgramDb};
53
54 type Events = Arc<Mutex<Vec<salsa::Event>>>;
55
56 #[salsa::db]
57 #[derive(Clone)]
58 pub(crate) struct TestDb {
59 storage: salsa::Storage<Self>,
60 files: Files,
61 system: TestSystem,
62 vendored: VendoredFileSystem,
63 program_settings: ProgramSettings,
64 }
65
66 impl TestDb {
67 fn new() -> Self {
68 let events = Events::default();
69 let vendored = ty_vendored::file_system().clone();
70 let program_settings = ProgramSettings::empty(&vendored);
71 Self {
72 storage: salsa::Storage::new(Some(Box::new({
73 move |event| {
74 tracing::trace!("event: {event:?}");
75 let mut events = events.lock().unwrap();
76 events.push(event);
77 }
78 }))),
79 system: TestSystem::default(),
80 vendored,
81 files: Files::default(),
82 program_settings,
83 }
84 }
85 }
86
87 impl DbWithTestSystem for TestDb {
88 fn test_system(&self) -> &TestSystem {
89 &self.system
90 }
91
92 fn test_system_mut(&mut self) -> &mut TestSystem {
93 &mut self.system
94 }
95 }
96
97 #[salsa::db]
98 impl SourceDb for TestDb {
99 fn vendored(&self) -> &VendoredFileSystem {
100 &self.vendored
101 }
102
103 fn system(&self) -> &dyn System {
104 &self.system
105 }
106
107 fn files(&self) -> &Files {
108 &self.files
109 }
110 }
111
112 #[salsa::db]
113 impl Db for TestDb {
114 fn should_check_file(&self, file: File) -> bool {
115 !file.path(self).is_vendored_path()
116 }
117 }
118
119 #[salsa::db]
120 impl ModuleResolverDb for TestDb {}
121
122 #[salsa::db]
123 impl TestProgramDb for TestDb {
124 fn program_settings(&self) -> &ProgramSettings {
125 &self.program_settings
126 }
127 }
128
129 #[salsa::db]
130 impl salsa::Database for TestDb {}
131
132 pub(crate) struct TestDbBuilder<'a> {
133 python_version: PythonVersion,
135 python_platform: PythonPlatform,
137 files: Vec<(&'a str, &'a str)>,
139 }
140
141 impl<'a> TestDbBuilder<'a> {
142 pub(crate) fn new() -> Self {
143 Self {
144 python_version: PythonVersion::default(),
145 python_platform: PythonPlatform::default(),
146 files: vec![],
147 }
148 }
149
150 pub(crate) fn with_file(
151 mut self,
152 path: &'a (impl AsRef<SystemPath> + ?Sized),
153 content: &'a str,
154 ) -> Self {
155 self.files.push((path.as_ref().as_str(), content));
156 self
157 }
158
159 pub(crate) fn build(self) -> anyhow::Result<TestDb> {
160 let mut db = TestDb::new();
161
162 let src_root = SystemPathBuf::from("/src");
163 db.memory_file_system().create_directory_all(&src_root)?;
164
165 db.write_files(self.files)
166 .context("Failed to write test files")?;
167
168 let program_settings = ProgramSettings {
169 python_version: PythonVersionWithSource {
170 version: self.python_version,
171 source: PythonVersionSource::default(),
172 },
173 python_platform: self.python_platform,
174 search_paths: SearchPathSettings::new(vec![src_root])
175 .to_search_paths(db.system(), db.vendored(), &FallibleStrategy)
176 .context("Invalid search path settings")?,
177 };
178 program_settings.search_paths.try_register_static_roots(&db);
179 db.program_settings = program_settings;
180
181 Ok(db)
182 }
183 }
184}