1use std::{collections::HashMap, io};
2
3use crate::{helpers::suffix_preference, providers::*, PythonVersion};
4use fancy_regex::Regex;
5use lazy_static::lazy_static;
6
7#[cfg(feature = "pyo3")]
8use pyo3::prelude::*;
9
10#[cfg(feature = "pyo3")]
11use crate::providers::pyobject::PyObjectProvider;
12
13lazy_static! {
14 static ref VERSION_REGEX: Regex = Regex::new(
15 r#"(?x)
16 ^(?P<major>\d+)(?:\.(?P<minor>\d+)(?:\.(?P<patch>[0-9]+))?)?\.?
17 (?:(?P<prerel>[abc]|rc|dev)(?:(?P<prerelversion>\d+(?:\.\d+)*))?)
18 ?(?P<postdev>(\.post(?P<post>\d+))?(\.dev(?P<dev>\d+))?)?
19 (?:-(?P<architecture>32|64))?"#
20 )
21 .unwrap();
22}
23
24#[cfg_attr(feature = "pyo3", pyclass)]
25pub struct Finder {
26 providers: Vec<Box<dyn Provider>>,
27 resolve_symlinks: bool,
28 same_file: bool,
29 same_interpreter: bool,
30}
31
32impl Default for Finder {
33 fn default() -> Self {
34 let f = Self {
35 providers: vec![],
36 resolve_symlinks: false,
37 same_file: true,
38 same_interpreter: true,
39 };
40 f.select_providers(&ALL_PROVIDERS[..]).unwrap()
41 }
42}
43
44impl Finder {
45 pub fn select_providers(mut self, names: &[&str]) -> Result<Self, io::Error> {
46 self.providers = names.iter().filter_map(|n| get_provider(*n)).collect();
47 Ok(self)
48 }
49
50 pub fn resolve_symlinks(mut self, resolve_symlinks: bool) -> Self {
51 self.resolve_symlinks = resolve_symlinks;
52 self
53 }
54
55 pub fn same_file(mut self, same_file: bool) -> Self {
56 self.same_file = same_file;
57 self
58 }
59
60 pub fn same_interpreter(mut self, same_interpreter: bool) -> Self {
61 self.same_interpreter = same_interpreter;
62 self
63 }
64
65 fn find_all_python_versions(&self) -> Vec<PythonVersion> {
66 self.providers
67 .iter()
68 .flat_map(|p| p.find_pythons())
69 .collect()
70 }
71
72 pub fn find_all(&self, options: MatchOptions) -> Vec<PythonVersion> {
73 let pythons = self.find_all_python_versions();
74 let mut filtered = vec![];
75 for python in pythons {
76 if python.matches(&options) {
77 filtered.push(python);
78 }
79 }
80 self.deduplicate(filtered)
81 }
82
83 pub fn find(&self, options: MatchOptions) -> Option<PythonVersion> {
84 self.find_all(options).first().cloned()
85 }
86
87 fn deduplicate_key(&self, python: &mut PythonVersion) -> String {
88 if !self.same_interpreter {
89 return python.interpreter().unwrap().to_str().unwrap().to_string();
90 }
91 if !self.same_file {
92 return python.content_hash().unwrap();
93 }
94 if self.resolve_symlinks && !python.keep_symlink {
95 return python.real_path().to_str().unwrap().to_string();
96 }
97 python.executable.to_str().unwrap().to_string()
98 }
99
100 fn deduplicate(&self, versions: Vec<PythonVersion>) -> Vec<PythonVersion> {
101 let mut result = HashMap::new();
102 let mut versions = versions;
103
104 versions.sort_by_cached_key(|p| {
105 (
106 p.executable.is_symlink(),
107 suffix_preference(&p.executable),
108 -(p.executable.to_string_lossy().len() as isize),
109 )
110 });
111
112 for version in versions.iter_mut() {
113 let key = self.deduplicate_key(version);
114 result.entry(key).or_insert(version.to_owned());
115 }
116 let mut py_versions = result.into_values().collect::<Vec<_>>();
117 py_versions.sort_by(|a, b| {
118 (b.version().unwrap(), b.executable.to_string_lossy().len())
119 .cmp(&(a.version().unwrap(), a.executable.to_string_lossy().len()))
120 });
121 py_versions
122 }
123}
124
125#[cfg(feature = "pyo3")]
126#[derive(FromPyObject)]
127pub enum StringInt {
128 STRING(String),
129 INT(usize),
130}
131
132#[cfg(feature = "pyo3")]
133#[pymethods]
134impl Finder {
135 #[new]
136 #[pyo3(signature = (resolve_symlinks = false, no_same_file = false, no_same_interpreter = false, selected_providers = None))]
137 fn py_new(
138 resolve_symlinks: bool,
139 no_same_file: bool,
140 no_same_interpreter: bool,
141 selected_providers: Option<Vec<String>>,
142 ) -> Result<Self, io::Error> {
143 let mut f = Self {
144 resolve_symlinks,
145 same_file: !no_same_file,
146 same_interpreter: !no_same_interpreter,
147 ..Self::default()
148 };
149 if let Some(names) = selected_providers {
150 let names: Vec<&str> = names.iter().map(|v| v.as_str()).collect();
151 f = f.select_providers(names.as_slice())?
152 }
153 Ok(f)
154 }
155
156 fn add_provider(&mut self, provider: PyObject, pos: Option<usize>) {
157 let provider = PyObjectProvider::new(provider);
158 if let Some(pos) = pos {
159 self.providers.insert(pos, Box::new(provider));
160 } else {
161 self.providers.push(Box::new(provider));
162 }
163 }
164
165 #[getter]
166 fn get_resolve_symlinks(&self) -> bool {
167 self.resolve_symlinks
168 }
169
170 #[setter]
171 fn set_resolve_symlinks(&mut self, resolve_symlinks: bool) {
172 self.resolve_symlinks = resolve_symlinks
173 }
174
175 #[getter]
176 fn get_same_file(&self) -> bool {
177 self.same_file
178 }
179
180 #[setter]
181 fn set_same_file(&mut self, same_file: bool) {
182 self.same_file = same_file
183 }
184
185 #[getter]
186 fn get_same_interpreter(&self) -> bool {
187 self.same_interpreter
188 }
189
190 #[setter]
191 fn set_same_interpreter(&mut self, same_interpreter: bool) {
192 self.same_interpreter = same_interpreter
193 }
194
195 #[pyo3(name = "find_all")]
196 pub fn py_find_all(
197 &self,
198 major: Option<StringInt>,
199 minor: Option<usize>,
200 patch: Option<usize>,
201 pre: Option<bool>,
202 dev: Option<bool>,
203 name: Option<String>,
204 architecture: Option<String>,
205 ) -> Vec<PythonVersion> {
206 let options = if let Some(StringInt::STRING(s)) = &major {
207 MatchOptions::default().version_spec(s)
208 } else {
209 MatchOptions {
210 major: if let Some(StringInt::INT(i)) = major {
211 Some(i)
212 } else {
213 None
214 },
215 minor,
216 patch,
217 pre,
218 dev,
219 name,
220 architecture,
221 }
222 };
223 self.find_all(options)
224 }
225
226 #[pyo3(name = "find")]
227 pub fn py_find(
228 &self,
229 major: Option<StringInt>,
230 minor: Option<usize>,
231 patch: Option<usize>,
232 pre: Option<bool>,
233 dev: Option<bool>,
234 name: Option<String>,
235 architecture: Option<String>,
236 ) -> Option<PythonVersion> {
237 self.py_find_all(major, minor, patch, pre, dev, name, architecture)
238 .first()
239 .cloned()
240 }
241}
242
243#[derive(Debug, Clone, PartialEq, Eq, Default)]
244pub struct MatchOptions {
245 pub major: Option<usize>,
246 pub minor: Option<usize>,
247 pub patch: Option<usize>,
248 pub pre: Option<bool>,
249 pub dev: Option<bool>,
250 pub name: Option<String>,
251 pub architecture: Option<String>,
252}
253
254impl MatchOptions {
255 fn from_version(version: &str) -> Option<Self> {
256 match VERSION_REGEX.captures(version) {
257 Ok(Some(capture)) => Some(Self {
258 major: capture.name("major").map(|m| m.as_str().parse().unwrap()),
259 minor: capture.name("minor").map(|m| m.as_str().parse().unwrap()),
260 patch: capture.name("patch").map(|m| m.as_str().parse().unwrap()),
261 pre: capture.name("prerel").map(|_| true),
262 dev: capture.name("dev").map(|_| true),
263 name: None,
264 architecture: capture
265 .name("architecture")
266 .map(|m| format!("{}bit", m.as_str())),
267 }),
268 _ => None,
269 }
270 }
271
272 pub fn version_spec(self, version: &str) -> Self {
273 if let Some(res) = Self::from_version(version) {
274 res
275 } else {
276 self.name(version)
277 }
278 }
279
280 pub fn major(mut self, major: usize) -> Self {
281 self.major = Some(major);
282 self
283 }
284
285 pub fn minor(mut self, minor: usize) -> Self {
286 self.minor = Some(minor);
287 self
288 }
289
290 pub fn patch(mut self, patch: usize) -> Self {
291 self.patch = Some(patch);
292 self
293 }
294
295 pub fn pre(mut self, pre: bool) -> Self {
296 self.pre = Some(pre);
297 self
298 }
299
300 pub fn dev(mut self, dev: bool) -> Self {
301 self.dev = Some(dev);
302 self
303 }
304
305 pub fn name(mut self, name: &str) -> Self {
306 self.name = Some(name.to_string());
307 self
308 }
309
310 pub fn architecture(mut self, architecture: &str) -> Self {
311 self.architecture = Some(architecture.to_string());
312 self
313 }
314}
315
316#[cfg(test)]
317mod test {
318 use super::*;
319
320 #[test]
321 fn test_find_pythons() {
322 let finder = Finder::default();
323
324 let pythons = finder.find_all(MatchOptions::default());
325 assert!(pythons.len() > 0);
326 }
327}