4 files changed, 95 insertions, 28 deletions

diff --git a/ci/setup-test-env.yml b/ci/setup-test-env.yml
index 80a031c1a..e36bca304 100644
--- a/ci/setup-test-env.yml
+++ b/ci/setup-test-env.yml
@@ -10,27 +10,16 @@ steps:
       versionSpec: "1.10"
     displayName: "Install a fixed version of Go"
 
-  # Because the Pipelines agent updates are out of sync, we can't install 3.6.9
-  # with PythonTool on macOS
-
+  # We are using pyenv to install Python for integration tests
   # Install Python
-  - task: UsePythonVersion@0
-    inputs:
-      versionSpec: "3.6.9"
-    displayName: "Install a fixed version of Python"
-    condition: not(eq(variables['Agent.OS'], 'Darwin'))
-  # Install Python (macOS)
   - script: |
       echo "##vso[task.setvariable variable=PYTHON_VERSION;]3.6.9"
       echo "##vso[task.setvariable variable=PYENV_ROOT;]$HOME/.pyenv"
-    condition: eq(variables['Agent.OS'], 'Darwin')
   - script: |
       curl https://pyenv.run | bash
       echo "##vso[task.setvariable variable=PATH;]$PYENV_ROOT/bin:$PYENV_ROOT/shims:$PATH"
-    condition: eq(variables['Agent.OS'], 'Darwin')
   - script: |
       eval "$(pyenv init -)"
       pyenv install $PYTHON_VERSION
       pyenv global $PYTHON_VERSION
-    condition: eq(variables['Agent.OS'], 'Darwin')
     displayName: "Install a fixed version of Python"
diff --git a/docs/config/README.md b/docs/config/README.md
index 8177cec59..1f32e2350 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -350,8 +350,13 @@ symbol = "🎁 "
 ## Python
 
 The `python` module shows the currently installed version of Python.
-It will also show the current Python virtual environment if one is
+
+If `pyenv_version_name` is set to `true`, it will display the pyenv version name.
+
+Otherwise, it will display the version number from `python --version`
+and show the current Python virtual environment if one is
 activated.
+
 The module will be shown if any of the following conditions are met:
 
 - The current directory contains a `.python-version` file
@@ -365,6 +370,9 @@ The module will be shown if any of the following conditions are met:
 | ---------- | ------- | -------------------------------------------------------- |
 | `symbol`   | `"🐍 "` | The symbol used before displaying the version of Python. |
 | `disabled` | `false` | Disables the `python` module.                            |
+| `pyenv_version_name` | `false` | Use pyenv to get Python version                            |
+| `pyenv_prefix` | `"pyenv "` | Prefix before pyenv version display (default display is `pyenv MY_VERSION`) |
+
 
 ### Example
 
@@ -373,6 +381,8 @@ The module will be shown if any of the following conditions are met:
 
 [python]
 symbol = "👾 "
+pyenv_version_name = true
+pyenv_prefix = "foo "
 ```
 
 ## Rust
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")
     }
+
 }
diff --git a/tests/testsuite/python.rs b/tests/testsuite/python.rs
index 9be7c9c0e..66515b43b 100644
--- a/tests/testsuite/python.rs
+++ b/tests/testsuite/python.rs
@@ -1,10 +1,10 @@
-use std::env;
 use std::fs::File;
 use std::io;
 
 use ansi_term::Color;
 
 use crate::common;
+use crate::common::TestCommand;
 
 #[test]
 #[ignore]
@@ -93,3 +93,42 @@ fn with_virtual_env() -> io::Result<()> {
     assert_eq!(expected, actual);
     Ok(())
 }
+
+#[test]
+#[ignore]
+fn with_pyenv() -> io::Result<()> {
+    let dir = common::new_tempdir()?;
+    File::create(dir.path().join("main.py"))?;
+    let output = common::render_module("python")
+        .use_config(toml::toml! {
+            [python]
+            pyenv_version_name = true
+        })
+        .env("VIRTUAL_ENV", "/foo/bar/my_venv")
+        .arg("--path")
+        .arg(dir.path())
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+    let expected = format!("via {} ", Color::Yellow.bold().paint("🐍 pyenv system"));
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
+#[test]
+#[ignore]
+fn with_pyenv_no_output() -> io::Result<()> {
+    let dir = common::new_tempdir()?;
+    File::create(dir.path().join("main.py"))?;
+    let output = common::render_module("python")
+        .use_config(toml::toml! {
+            [python]
+            pyenv_version_name = true
+        })
+        .env("PATH", "")
+        .arg("--path")
+        .arg(dir.path())
+        .output()?;
+    let actual = String::from_utf8(output.stdout).unwrap();
+    assert_eq!("", actual);
+    Ok(())
+}