Skip to main content

ty_module_resolver/
db.rs

1use ruff_db::Db as SourceDb;
2
3#[salsa::db]
4pub trait Db: SourceDb {}
5
6#[cfg(test)]
7pub(crate) mod tests {
8    use std::sync::{Arc, Mutex};
9
10    use ruff_db::Db as SourceDb;
11    use ruff_db::files::Files;
12    use ruff_db::system::{DbWithTestSystem, TestSystem};
13    use ruff_db::vendored::VendoredFileSystem;
14    use ruff_python_ast::PythonVersion;
15
16    use super::Db;
17    use crate::{ResolverEnvironment, resolve::SearchPaths};
18
19    type Events = Arc<Mutex<Vec<salsa::Event>>>;
20
21    #[salsa::db]
22    #[derive(Clone)]
23    pub(crate) struct TestDb {
24        storage: salsa::Storage<Self>,
25        files: Files,
26        system: TestSystem,
27        vendored: VendoredFileSystem,
28        search_paths: Arc<SearchPaths>,
29        python_version: PythonVersion,
30        events: Events,
31    }
32
33    impl TestDb {
34        pub(crate) fn new() -> Self {
35            let events = Events::default();
36            Self {
37                storage: salsa::Storage::new(Some(Box::new({
38                    let events = events.clone();
39                    move |event| {
40                        tracing::trace!("event: {event:?}");
41                        let mut events = events.lock().unwrap();
42                        events.push(event);
43                    }
44                }))),
45                system: TestSystem::default(),
46                vendored: ty_vendored::file_system().clone(),
47                files: Files::default(),
48                search_paths: Arc::new(SearchPaths::empty(ty_vendored::file_system())),
49                python_version: PythonVersion::default(),
50                events,
51            }
52        }
53
54        pub(crate) fn with_search_paths(mut self, search_paths: SearchPaths) -> Self {
55            self.set_search_paths(search_paths);
56            self
57        }
58
59        pub(crate) fn with_python_version(mut self, python_version: PythonVersion) -> Self {
60            self.python_version = python_version;
61            self
62        }
63
64        pub(crate) fn set_search_paths(&mut self, search_paths: SearchPaths) {
65            search_paths.try_register_static_roots(self);
66            self.search_paths = Arc::new(search_paths);
67        }
68
69        pub(crate) fn search_paths(&self) -> &SearchPaths {
70            &self.search_paths
71        }
72
73        pub(crate) fn resolver_environment(&self) -> ResolverEnvironment<'_> {
74            ResolverEnvironment::new(self, self.python_version, self.search_paths.as_ref())
75        }
76
77        /// Takes the salsa events.
78        pub(crate) fn take_salsa_events(&mut self) -> Vec<salsa::Event> {
79            let mut events = self.events.lock().unwrap();
80            std::mem::take(&mut *events)
81        }
82
83        /// Clears the salsa events.
84        pub(crate) fn clear_salsa_events(&mut self) {
85            self.take_salsa_events();
86        }
87    }
88
89    impl DbWithTestSystem for TestDb {
90        fn test_system(&self) -> &TestSystem {
91            &self.system
92        }
93
94        fn test_system_mut(&mut self) -> &mut TestSystem {
95            &mut self.system
96        }
97    }
98
99    #[salsa::db]
100    impl SourceDb for TestDb {
101        fn vendored(&self) -> &VendoredFileSystem {
102            &self.vendored
103        }
104
105        fn system(&self) -> &dyn ruff_db::system::System {
106            &self.system
107        }
108
109        fn files(&self) -> &Files {
110            &self.files
111        }
112    }
113
114    #[salsa::db]
115    impl Db for TestDb {}
116
117    #[salsa::db]
118    impl salsa::Database for TestDb {}
119}