From ca12d22acea70cb3fb3b05da16388f34f22b8fb8 Mon Sep 17 00:00:00 2001
From: John Letey <john.letey@colorado.edu>
Date: Thu, 25 Apr 2019 16:06:18 +0100
Subject: [PATCH] Add Python version segment (#23)

---
 README.md             |  3 +-
 src/modules/mod.rs    |  2 ++
 src/modules/python.rs | 76 +++++++++++++++++++++++++++++++++++++++++++
 src/print.rs          |  9 ++++-
 4 files changed, 88 insertions(+), 2 deletions(-)
 create mode 100644 src/modules/python.rs

diff --git a/README.md b/README.md
index 91c30d48..926a2064 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,7 @@ I'm very new to Rust, so any help is appreciated when it comes to improving deve
 - [x] Prompt character turns red if the last command exits with non-zero code.
 - [x] Current Node.js version(`⬢`).
 - [x] Current Rust version (`🦀`).
+- [x] Current Python version (`🐍`).
 - [ ] Package version of package in current directory (`📦`).
 - [ ] Current battery level and status
 - [ ] Current Git branch and rich repo status.
@@ -105,4 +106,4 @@ To test locally run the below command:
 
 ```bash
 cargo run -- $status
-```
\ No newline at end of file
+```
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 10f1c654..8e034331 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -2,6 +2,7 @@ mod character;
 mod directory;
 mod line_break;
 mod nodejs;
+mod python;
 mod rust;
 
 use crate::context::Context;
@@ -13,6 +14,7 @@ pub fn handle(module: &str, context: &Context) -> Option<Segment> {
         "char" | "character" => character::segment(context),
         "node" | "nodejs" => nodejs::segment(context),
         "rust" | "rustlang" => rust::segment(context),
+        "python" => python::segment(context),
         "line_break" => line_break::segment(context),
 
         _ => panic!("Unknown module: {}", module),
diff --git a/src/modules/python.rs b/src/modules/python.rs
new file mode 100644
index 00000000..3d182bb2
--- /dev/null
+++ b/src/modules/python.rs
@@ -0,0 +1,76 @@
+use super::Segment;
+use crate::context::Context;
+use ansi_term::Color;
+use std::path::PathBuf;
+use std::process::Command;
+
+/// Creates a segment with the current Python version
+///
+/// Will display the Python version if any of the following criteria are met:
+///     - Current directory contains a `.py` file
+///     - Current directory contains a `.python-version` file
+///     - Current directory contains a `requirements.txt` file
+///     - Current directory contains a `pyproject.toml` file
+pub fn segment(context: &Context) -> Option<Segment> {
+    let is_py_project = context.dir_files.iter().any(has_py_files);
+    if !is_py_project {
+        return None;
+    }
+
+    match get_python_version() {
+        Some(python_version) => {
+            const PYTHON_CHAR: &str = "🐍";
+            const SEGMENT_COLOR: Color = Color::Yellow;
+
+            let mut segment = Segment::new("python");
+            segment.set_style(SEGMENT_COLOR);
+
+            let formatted_version = format_python_version(python_version);
+            segment.set_value(format!("{} {}", PYTHON_CHAR, formatted_version));
+
+            Some(segment)
+        }
+        None => None,
+    }
+}
+
+fn has_py_files(dir_entry: &PathBuf) -> bool {
+    let is_py_file =
+        |d: &PathBuf| -> bool { d.is_file() && d.extension().unwrap_or_default() == "py" };
+    let is_python_version = |d: &PathBuf| -> bool {
+        d.is_file() && d.file_name().unwrap_or_default() == ".python-version"
+    };
+    let is_requirements_txt = |d: &PathBuf| -> bool {
+        d.is_file() && d.file_name().unwrap_or_default() == "requirements.txt"
+    };
+    let is_py_project = |d: &PathBuf| -> bool {
+        d.is_file() && d.file_name().unwrap_or_default() == "pyproject.toml"
+    };
+
+    is_py_file(&dir_entry)
+        || is_python_version(&dir_entry)
+        || is_requirements_txt(&dir_entry)
+        || is_py_project(&dir_entry)
+}
+
+fn get_python_version() -> Option<String> {
+    match Command::new("python").arg("--version").output() {
+        Ok(output) => Some(String::from_utf8(output.stdout).unwrap()),
+        Err(_) => None,
+    }
+}
+
+fn format_python_version(python_stdout: String) -> String {
+    format!("v{}", python_stdout.trim_start_matches("Python ").trim())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_format_python_version() {
+        let input = String::from("Python 3.7.2");
+        assert_eq!(format_python_version(input), "v3.7.2");
+    }
+}
diff --git a/src/print.rs b/src/print.rs
index 5804e0e7..836795d4 100644
--- a/src/print.rs
+++ b/src/print.rs
@@ -5,7 +5,14 @@ use crate::context::Context;
 use crate::modules;
 
 pub fn prompt(args: ArgMatches) {
-    let prompt_order = vec!["directory", "nodejs", "rust", "line_break", "character"];
+    let prompt_order = vec![
+        "directory",
+        "nodejs",
+        "rust",
+        "python",
+        "line_break",
+        "character",
+    ];
     let context = Context::new(args);
 
     // TODO: