findpython/
finder.rs

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}