use std::{
collections, env,
io::{self, BufRead},
path,
str::FromStr,
};
type VersionComponent = u16;
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum RequestedVersion {
Any,
Loose(VersionComponent),
Exact(VersionComponent, VersionComponent),
}
impl FromStr for RequestedVersion {
type Err = String;
fn from_str(version_string: &str) -> Result<Self, Self::Err> {
if version_string.is_empty() {
return Err("version string is empty".to_string());
}
let mut char_iter = version_string.chars();
let mut major_ver = Vec::new();
let mut dot = false;
for c in char_iter.by_ref() {
if c == '.' {
dot = true;
break;
} else if c.is_ascii_digit() {
major_ver.push(c);
} else {
return Err(format!(
"{:?} contains a non-numeric and non-period character",
version_string
));
}
}
let mut minor_ver = Vec::new();
if dot {
for c in char_iter.by_ref() {
if c.is_ascii_digit() {
minor_ver.push(c);
} else {
return Err(format!(
"{:?} contains a non-numeric character after a period",
version_string
));
}
}
}
let major = char_vec_to_int(&major_ver)?;
if !dot {
Ok(RequestedVersion::Loose(major))
} else if minor_ver.is_empty() {
Err(format!(
"{:?} is missing a minor version number",
version_string
))
} else {
let minor = char_vec_to_int(&minor_ver)?;
Ok(RequestedVersion::Exact(major, minor))
}
}
}
impl RequestedVersion {
pub fn env_var(&self) -> Option<String> {
match self {
RequestedVersion::Any => Some("PY_PYTHON".to_string()),
RequestedVersion::Loose(major) => Some(format!("PY_PYTHON{}", major)),
_ => None,
}
}
}
#[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Version {
major: VersionComponent,
minor: VersionComponent,
}
#[derive(Debug, PartialEq)]
pub enum VersionMatch {
NotAtAll,
Loosely,
Exactly,
}
impl Version {
pub fn matches(&self, requested: RequestedVersion) -> VersionMatch {
match requested {
RequestedVersion::Any => VersionMatch::Loosely,
RequestedVersion::Loose(major) => {
if self.major == major {
VersionMatch::Loosely
} else {
VersionMatch::NotAtAll
}
}
RequestedVersion::Exact(major, minor) => {
if self.major == major && self.minor == minor {
VersionMatch::Exactly
} else {
VersionMatch::NotAtAll
}
}
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum Action {
Help(path::PathBuf),
Execute {
launcher: path::PathBuf,
version: RequestedVersion,
args: Vec<String>,
},
}
pub fn action_from_args(mut args: Vec<String>) -> Action {
let launcher_path = path::PathBuf::from(args.remove(0));
if !args.is_empty() {
let flag = &args[0];
if flag == "-h" || flag == "--help" {
return Action::Help(launcher_path);
} else if let Some(version) = version_from_flag(&flag) {
args.remove(0);
return Action::Execute {
launcher: launcher_path,
version,
args,
};
}
}
Action::Execute {
launcher: launcher_path,
version: RequestedVersion::Any,
args,
}
}
fn char_vec_to_int(char_vec: &[char]) -> Result<VersionComponent, String> {
let joined_string = char_vec.iter().collect::<String>();
let parse_result = joined_string.parse();
parse_result.or_else(|_| Err(format!("error converting {:?} to a number", joined_string)))
}
pub fn version_from_flag(arg: &str) -> Option<RequestedVersion> {
if !arg.starts_with('-') {
None
} else {
RequestedVersion::from_str(&arg[1..]).ok()
}
}
pub fn virtual_env() -> Option<path::PathBuf> {
match env::var_os("VIRTUAL_ENV") {
None => None,
Some(venv_root) => {
let mut path = path::PathBuf::new();
path.push(venv_root);
path.push("bin");
path.push("python");
Some(path)
}
}
}
pub fn path_entries() -> Vec<path::PathBuf> {
if let Some(path_val) = env::var_os("PATH") {
env::split_paths(&path_val).collect()
} else {
Vec::new()
}
}
pub fn directory_contents(path: &path::Path) -> collections::HashSet<path::PathBuf> {
let mut paths = collections::HashSet::new();
if let Ok(contents) = path.read_dir() {
for content in contents {
if let Ok(found_content) = content {
paths.insert(found_content.path());
}
}
}
paths
}
pub fn filter_python_executables(
paths: collections::HashSet<path::PathBuf>,
) -> collections::HashMap<Version, path::PathBuf> {
let mut executables = collections::HashMap::new();
for path in paths {
let unencoded_file_name = match path.file_name() {
Some(x) => x,
None => continue,
};
let file_name = match unencoded_file_name.to_str() {
Some(x) => x,
None => continue,
};
if file_name.len() < "python3.0".len() || !file_name.starts_with("python") {
continue;
}
let version_part = &file_name["python".len()..];
if let Ok(found_version) = RequestedVersion::from_str(&version_part) {
match found_version {
RequestedVersion::Exact(major, minor) => {
executables.insert(Version { major, minor }, path.clone())
}
_ => continue,
};
}
}
executables
}
pub fn available_executables(
requested_version: RequestedVersion,
) -> collections::HashMap<Version, path::PathBuf> {
let mut found_versions = collections::HashMap::new();
for path in path_entries() {
let all_contents = directory_contents(&path);
for (version, path) in filter_python_executables(all_contents) {
match version.matches(requested_version) {
VersionMatch::NotAtAll => continue,
VersionMatch::Loosely => {
if !found_versions.contains_key(&version) && path.is_file() {
found_versions.insert(version, path);
}
}
VersionMatch::Exactly => {
if path.is_file() {
found_versions.insert(version, path);
break;
}
}
};
}
}
found_versions
}
pub fn choose_executable(
version_paths: &collections::HashMap<Version, path::PathBuf>,
) -> Option<path::PathBuf> {
let mut pairs: Vec<(&Version, &path::PathBuf)> = version_paths.iter().collect();
pairs.sort_unstable_by_key(|p| p.0);
pairs.last().map(|(_, path)| path.to_path_buf())
}
pub fn find_shebang(reader: impl io::Read) -> Option<String> {
let mut buffered_reader = io::BufReader::new(reader);
let mut line = String::new();
if buffered_reader.read_line(&mut line).is_err() {
return None;
};
if !line.starts_with("#!") {
None
} else {
Some(line[2..].trim().to_string())
}
}
pub fn split_shebang(shebang_line: &str) -> Option<(RequestedVersion, Vec<String>)> {
let accepted_paths = [
"/usr/bin/python",
"/usr/local/bin/python",
"/usr/bin/env python",
"python",
];
for exec_path in &accepted_paths {
if !shebang_line.starts_with(exec_path) {
continue;
}
let trimmed_shebang = shebang_line[exec_path.len()..].to_string();
let version_string: String = trimmed_shebang
.chars()
.take_while(|c| c.is_ascii_digit() || *c == '.')
.collect();
let specified_version = if version_string.is_empty() {
Ok(RequestedVersion::Loose(2))
} else {
RequestedVersion::from_str(&version_string)
};
return specified_version
.map(|version| {
let args = trimmed_shebang[version_string.len()..].trim();
(
version,
args.split_whitespace().map(|s| s.to_string()).collect(),
)
})
.ok();
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[allow(non_snake_case)]
fn test_RequestedVersion_from_string() {
assert!(RequestedVersion::from_str(&".3".to_string()).is_err());
assert!(RequestedVersion::from_str(&"3.".to_string()).is_err());
assert!(RequestedVersion::from_str(&"h".to_string()).is_err());
assert!(RequestedVersion::from_str(&"3.b".to_string()).is_err());
assert!(RequestedVersion::from_str(&"a.7".to_string()).is_err());
assert_eq!(
RequestedVersion::from_str(&"3".to_string()),
Ok(RequestedVersion::Loose(3))
);
assert_eq!(
RequestedVersion::from_str(&"3.8".to_string()),
Ok(RequestedVersion::Exact(3, 8))
);
assert_eq!(
RequestedVersion::from_str(&"42.13".to_string()),
Ok(RequestedVersion::Exact(42, 13))
);
assert!(RequestedVersion::from_str(&"3.6.5".to_string()).is_err());
}
#[test]
fn test_version_from_flag() {
assert!(version_from_flag(&"-S".to_string()).is_none());
assert!(version_from_flag(&"--something".to_string()).is_none());
assert_eq!(
version_from_flag(&"-3".to_string()),
Some(RequestedVersion::Loose(3))
);
assert_eq!(
version_from_flag(&"-3.6".to_string()),
Some(RequestedVersion::Exact(3, 6))
);
assert_eq!(
version_from_flag(&"-42.13".to_string()),
Some(RequestedVersion::Exact(42, 13))
);
assert!(version_from_flag(&"-3.6.4".to_string()).is_none());
}
#[test]
fn test_action_from_args() {
let launcher_string = String::from("/py");
let launcher_path = path::PathBuf::from(&launcher_string);
assert_eq!(
action_from_args(vec![launcher_string.clone()]),
Action::Execute {
launcher: launcher_path.clone(),
version: RequestedVersion::Any,
args: Vec::new(),
}
);
assert_eq!(
action_from_args(vec![launcher_string.clone(), String::from("-h")]),
Action::Help(launcher_path.clone())
);
assert_eq!(
action_from_args(vec![launcher_string.clone(), String::from("--help")]),
Action::Help(launcher_path.clone())
);
assert_eq!(
action_from_args(vec![launcher_string.clone(), String::from("-V")]),
Action::Execute {
launcher: launcher_path.clone(),
version: RequestedVersion::Any,
args: vec![String::from("-V")],
}
);
assert_eq!(
action_from_args(vec![launcher_string.clone(), String::from("-3")]),
Action::Execute {
launcher: launcher_path.clone(),
version: RequestedVersion::Loose(3),
args: Vec::new(),
}
);
assert_eq!(
action_from_args(vec![launcher_string.clone(), String::from("-3.6")]),
Action::Execute {
launcher: launcher_path.clone(),
version: RequestedVersion::Exact(3, 6),
args: Vec::new(),
}
);
assert_eq!(
action_from_args(vec![
launcher_string.clone(),
String::from("-3.6"),
String::from("script.py"),
]),
Action::Execute {
launcher: launcher_path.clone(),
version: RequestedVersion::Exact(3, 6),
args: vec![String::from("script.py")],
}
);
assert_eq!(
action_from_args(vec![launcher_string.clone(), String::from("script.py")]),
Action::Execute {
launcher: launcher_path.clone(),
version: RequestedVersion::Any,
args: vec![String::from("script.py")],
}
);
}
#[test]
fn test_virtual_env() {
let original_venv = env::var_os("VIRTUAL_ENV");
env::remove_var("VIRTUAL_ENV");
assert_eq!(virtual_env(), None);
env::set_var("VIRTUAL_ENV", "/some/path");
assert_eq!(
virtual_env(),
Some(path::PathBuf::from("/some/path/bin/python"))
);
match original_venv {
None => env::remove_var("VIRTUAL_ENV"),
Some(venv_value) => env::set_var("VIRTUAL_ENV", venv_value),
}
}
#[test]
fn unit_test_path_entries() {
let paths = vec!["/a", "/b", "/c"];
if let Ok(joined_paths) = env::join_paths(&paths) {
let original_paths = env::var_os("PATH");
env::set_var("PATH", joined_paths);
assert_eq!(
path_entries(),
paths
.iter()
.map(|p| path::PathBuf::from(p))
.collect::<Vec<path::PathBuf>>()
);
match original_paths {
Some(paths) => env::set_var("PATH", paths),
None => env::set_var("PATH", ""),
}
}
}
#[test]
fn system_test_path_entries() {
if let Some(paths) = env::var_os("PATH") {
let found_paths = path_entries();
assert_eq!(found_paths.len(), env::split_paths(&paths).count());
for (index, path) in env::split_paths(&paths).enumerate() {
assert_eq!(found_paths[index], path);
}
}
}
#[test]
fn test_filter_python_executables() {
let paths = vec![
"/bad/path/python",
"/bad/path/python3",
"/bad/path/hello",
"/bad/path/pytho3.6",
"/bad/path/rython3.6",
"/good/path/python3.6",
"/good/python42.13",
];
let all_paths = paths
.iter()
.map(path::PathBuf::from)
.collect::<collections::HashSet<path::PathBuf>>();
let results = filter_python_executables(all_paths);
let good_version1 = Version { major: 3, minor: 6 };
let good_version2 = Version {
major: 42,
minor: 13,
};
let mut expected = paths[5];
match results.get(&good_version1) {
Some(path) => assert_eq!(*path, path::PathBuf::from(expected)),
None => panic!("{:?} not found", good_version1),
};
expected = paths[6];
match results.get(&good_version2) {
Some(path) => assert_eq!(*path, path::PathBuf::from(expected)),
None => panic!("{:?} not found", good_version2),
}
}
#[test]
fn test_version_matches() {
let any = RequestedVersion::Any;
let loose_42 = RequestedVersion::Loose(42);
let exact_42_13 = RequestedVersion::Exact(42, 13);
let version_3_6 = Version { major: 3, minor: 6 };
let version_42_0 = Version {
major: 42,
minor: 0,
};
let version_42_13 = Version {
major: 42,
minor: 13,
};
assert_eq!(version_3_6.matches(any), VersionMatch::Loosely);
assert_eq!(version_3_6.matches(loose_42), VersionMatch::NotAtAll);
assert_eq!(version_3_6.matches(exact_42_13), VersionMatch::NotAtAll);
assert_eq!(version_42_0.matches(any), VersionMatch::Loosely);
assert_eq!(version_42_0.matches(loose_42), VersionMatch::Loosely);
assert_eq!(version_42_0.matches(exact_42_13), VersionMatch::NotAtAll);
assert_eq!(version_42_13.matches(any), VersionMatch::Loosely);
assert_eq!(version_42_13.matches(loose_42), VersionMatch::Loosely);
assert_eq!(version_42_13.matches(exact_42_13), VersionMatch::Exactly);
}
#[test]
fn test_choose_executable() {
let version_3_6 = Version { major: 3, minor: 6 };
let version_42_0 = Version {
major: 42,
minor: 0,
};
let version_42_13 = Version {
major: 42,
minor: 13,
};
let path_3_6 = path::PathBuf::from("/python3.6");
let path_42_0 = path::PathBuf::from("/python42.0");
let path_42_13 = path::PathBuf::from("/python42.13");
let mut mapping = collections::HashMap::new();
if choose_executable(&mapping).is_some() {
panic!("found a non-existent path");
};
mapping.insert(version_3_6, path_3_6.clone());
match choose_executable(&mapping) {
Some(path) => assert_eq!(path, path_3_6),
None => panic!("no path found"),
}
mapping.insert(version_42_0, path_42_0);
mapping.insert(version_42_13, path_42_13.clone());
match choose_executable(&mapping) {
Some(path) => assert_eq!(path, path_42_13),
None => panic!("no path found"),
}
}
#[test]
fn test_requested_version_env_var() {
assert_eq!(
RequestedVersion::Any.env_var(),
Some("PY_PYTHON".to_string())
);
assert_eq!(
RequestedVersion::Loose(3).env_var(),
Some("PY_PYTHON3".to_string())
);
assert_eq!(
RequestedVersion::Loose(42).env_var(),
Some("PY_PYTHON42".to_string())
);
assert!(RequestedVersion::Exact(3, 8).env_var().is_none());
}
#[test]
fn test_find_shebang() {
assert_eq!(
find_shebang("#! /usr/bin/cat\nprint('Hello!')\n".as_bytes()),
Some("/usr/bin/cat".to_string())
);
assert_eq!(find_shebang("print('Hello!')".as_bytes()), None);
assert_eq!(
find_shebang("#!/usr/bin/cat\nHello".as_bytes()),
Some("/usr/bin/cat".to_string())
);
assert_eq!(
find_shebang("#! /usr/bin/env python -S".as_bytes()),
Some("/usr/bin/env python -S".to_string())
);
assert_eq!(
find_shebang("#! /usr/bin/python \n# Hello".as_bytes()),
Some("/usr/bin/python".to_string())
);
assert_eq!(
find_shebang("#! /usr/bin/python".as_bytes()),
Some("/usr/bin/python".to_string())
);
}
#[test]
fn test_split_shebang() {
assert_eq!(split_shebang(&"/usr/bin/rustup".to_string()), None);
assert_eq!(
split_shebang(&"/usr/bin/rustup self update".to_string()),
None
);
assert_eq!(
split_shebang(&"/usr/bin/env python".to_string()),
Some((RequestedVersion::Loose(2), Vec::new()))
);
assert_eq!(
split_shebang(&"/usr/bin/python42.13".to_string()),
Some((RequestedVersion::Exact(42, 13), Vec::new()))
);
assert_eq!(
split_shebang(&"python -S -v".to_string()),
Some((
RequestedVersion::Loose(2),
vec!["-S".to_string(), "-v".to_string()]
))
);
assert_eq!(
split_shebang(&"/usr/local/bin/python3.7 -S".to_string()),
Some((RequestedVersion::Exact(3, 7), vec!["-S".to_string()]))
);
}
}