Skip to main content

ty_python_core/
db.rs

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/// Database giving access to semantic information about a Python program.
8#[salsa::db]
9pub trait Db: ModuleResolverDb {
10    /// Returns `true` if the file should be checked.
11    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    // Salsa-cached because interning a program requires hashing all search paths.
20    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        /// Target Python version
134        python_version: PythonVersion,
135        /// Target Python platform
136        python_platform: PythonPlatform,
137        /// Path and content pairs for files that should be present
138        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}