feat: Implement configuration to display pyenv version name (#140)

This behavior can be enabled via setting `use_pyenv` to true.
The "pyenv" prefix before the version name can be configured using `pyenv_prefix`.

1 files changed, 43 insertions, 14 deletions

diff --git a/src/modules/python.rs b/src/modules/python.rs
index 0e7b3c250..8c9cc33a7 100644
--- a/src/modules/python.rs
+++ b/src/modules/python.rs
@@ -4,6 +4,8 @@ use std::process::Command;
 
 use ansi_term::Color;
 
+use crate::config::Config;
+
 use super::{Context, Module};
 
 /// Creates a module with the current Python version
@@ -24,26 +26,52 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
         return None;
     }
 
-    match get_python_version() {
-        Some(python_version) => {
-            const PYTHON_CHAR: &str = "🐍 ";
-            let module_color = Color::Yellow.bold();
+    let mut module = context.new_module("python")?;
+    let pyenv_version_name = module
+        .config_value_bool("pyenv_version_name")
+        .unwrap_or(false);
 
-            let mut module = context.new_module("python")?;
-            module.set_style(module_color);
+    const PYTHON_CHAR: &str = "🐍 ";
+    let module_color = Color::Yellow.bold();
+    module.set_style(module_color);
+    module.new_segment("symbol", PYTHON_CHAR);
 
-            let formatted_version = format_python_version(&python_version);
-            module.new_segment("symbol", PYTHON_CHAR);
-            module.new_segment("version", &formatted_version);
-            get_python_virtual_env()
-                .map(|virtual_env| module.new_segment("virtualenv", &format!("({})", virtual_env)));
+    select_python_version(pyenv_version_name)
+        .map(|python_version| python_module(module, pyenv_version_name, python_version))
+}
 
-            Some(module)
-        }
-        None => None,
+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) => {
@@ -95,4 +123,5 @@ mod tests {
         env::set_var("VIRTUAL_ENV", "/foo/bar/my_venv");
         assert_eq!(get_python_virtual_env().unwrap(), "my_venv")
     }
+
 }