diff options
author | Matan Kushner <hello@matchai.me> | 2019-10-04 22:30:46 +0900 |
---|---|---|
committer | GitHub <noreply@github.com> | 2019-10-04 22:30:46 +0900 |
commit | 05210b9510b797f7738d5b2d51e8a6877f2d5283 (patch) | |
tree | 7399401dba9373f61035dbbd055f4137cd20f705 /src/modules/python.rs | |
parent | e90a3768da7882db092b38d141cf8e19fabbee56 (diff) |
refactor: Go from Rust workspaces to a package with nested packages (#480)
Diffstat (limited to 'src/modules/python.rs')
-rw-r--r-- | src/modules/python.rs | 122 |
1 files changed, 122 insertions, 0 deletions
diff --git a/src/modules/python.rs b/src/modules/python.rs new file mode 100644 index 000000000..b871ef808 --- /dev/null +++ b/src/modules/python.rs @@ -0,0 +1,122 @@ +use std::env; +use std::path::Path; +use std::process::Command; + +use ansi_term::Color; + +use super::{Context, Module}; + +/// Creates a module with the current Python version +/// +/// Will display the Python version if any of the following criteria are met: +/// - Current directory contains a `.python-version` file +/// - Current directory contains a `requirements.txt` file +/// - Current directory contains a `pyproject.toml` file +/// - Current directory contains a file with the `.py` extension +/// - Current directory contains a `Pipfile` file +/// - Current directory contains a `tox.ini` file +pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> { + let is_py_project = context + .try_begin_scan()? + .set_files(&[ + "requirements.txt", + ".python-version", + "pyproject.toml", + "Pipfile", + "tox.ini", + ]) + .set_extensions(&["py"]) + .is_match(); + + if !is_py_project { + return None; + } + + let mut module = context.new_module("python"); + let pyenv_version_name = module + .config_value_bool("pyenv_version_name") + .unwrap_or(false); + + const PYTHON_CHAR: &str = "🐍 "; + let module_color = module + .config_value_style("style") + .unwrap_or_else(|| Color::Yellow.bold()); + module.set_style(module_color); + module.new_segment("symbol", PYTHON_CHAR); + + select_python_version(pyenv_version_name) + .map(|python_version| python_module(module, pyenv_version_name, python_version)) +} + +fn python_module(mut module: Module, pyenv_version_name: bool, python_version: String) -> Module { + const PYENV_PREFIX: &str = "pyenv "; + + if pyenv_version_name { + module.new_segment("pyenv_prefix", PYENV_PREFIX); + module.new_segment("version", &python_version.trim()); + } else { + let formatted_version = format_python_version(&python_version); + module.new_segment("version", &formatted_version); + get_python_virtual_env() + .map(|virtual_env| module.new_segment("virtualenv", &format!("({})", virtual_env))); + }; + + module +} + +fn select_python_version(pyenv_version_name: bool) -> Option<String> { + if pyenv_version_name { + get_pyenv_version() + } else { + get_python_version() + } +} + +fn get_pyenv_version() -> Option<String> { + Command::new("pyenv") + .arg("version-name") + .output() + .ok() + .and_then(|output| String::from_utf8(output.stdout).ok()) +} + +fn get_python_version() -> Option<String> { + match Command::new("python").arg("--version").output() { + Ok(output) => { + // We have to check both stdout and stderr since for Python versions + // < 3.4, Python reports to stderr and for Python version >= 3.5, + // Python reports to stdout + if output.stdout.is_empty() { + let stderr_string = String::from_utf8(output.stderr).unwrap(); + Some(stderr_string) + } else { + let stdout_string = String::from_utf8(output.stdout).unwrap(); + Some(stdout_string) + } + } + Err(_) => None, + } +} + +fn format_python_version(python_stdout: &str) -> String { + format!("v{}", python_stdout.trim_start_matches("Python ").trim()) +} + +fn get_python_virtual_env() -> Option<String> { + env::var("VIRTUAL_ENV").ok().and_then(|venv| { + Path::new(&venv) + .file_name() + .map(|filename| String::from(filename.to_str().unwrap_or(""))) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_format_python_version() { + let input = "Python 3.7.2"; + assert_eq!(format_python_version(input), "v3.7.2"); + } +} |