From e262187f4ca276ddc9ae88ee19e82c4307ffb0ce Mon Sep 17 00:00:00 2001
From: Kevin Lane <lane.kevin.a@gmail.com>
Date: Thu, 15 Aug 2019 13:41:06 -0700
Subject: [PATCH] feat: Display Python package version from poetry (#153)

---
 docs/config/README.md  |  6 +++--
 src/modules/package.rs | 54 +++++++++++++++++++++++++++++++++++-------
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/docs/config/README.md b/docs/config/README.md
index 2f51455c..d6e3c2f9 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -348,13 +348,15 @@ symbol = "🤖 "
 ## Package Version
 
 The `package` module is shown when the current directory is the repository for a
-package, and shows its current version. The module currently supports `npm` and
-`cargo` packages.
+package, and shows its current version. The module currently supports `npm`, `cargo`,
+and `poetry` packages.
 
 - **npm** – The `npm` package version is extracted from the `package.json` present
   in the current directory
 - **cargo** – The `cargo` package version is extracted from the `Cargo.toml` present
   in the current directory
+- **poetry** – The `poetry` package version is extracted from the `pyproject.toml` present
+  in the current directory
 
 > ⚠️ The version being shown is that of the package whose source code is in your
 > current directory, not your package manager.
diff --git a/src/modules/package.rs b/src/modules/package.rs
index 044499f5..5edacd21 100644
--- a/src/modules/package.rs
+++ b/src/modules/package.rs
@@ -46,18 +46,28 @@ fn extract_package_version(file_contents: &str) -> Option<String> {
     Some(formatted_version)
 }
 
+fn extract_poetry_version(file_contents: &str) -> Option<String> {
+    let poetry_toml: toml::Value = toml::from_str(file_contents).ok()?;
+    let raw_version = poetry_toml
+        .get("tool")?
+        .get("poetry")?
+        .get("version")?
+        .as_str()?;
+
+    let formatted_version = format_version(raw_version);
+    Some(formatted_version)
+}
+
 fn get_package_version() -> Option<String> {
-    let cargo_toml = utils::read_file("Cargo.toml");
-    if let Ok(cargo_toml) = cargo_toml {
+    if let Ok(cargo_toml) = utils::read_file("Cargo.toml") {
         return extract_cargo_version(&cargo_toml);
-    }
-
-    let package_json = utils::read_file("package.json");
-    if let Ok(package_json) = package_json {
+    } else if let Ok(package_json) = utils::read_file("package.json") {
         return extract_package_version(&package_json);
+    } else if let Ok(poetry_toml) = utils::read_file("pyproject.toml") {
+        return extract_poetry_version(&poetry_toml);
+    } else {
+        None
     }
-
-    None
 }
 
 fn format_version(version: &str) -> String {
@@ -123,4 +133,32 @@ mod tests {
             expected_version
         );
     }
+
+    #[test]
+    fn test_extract_poetry_version() {
+        let poetry_with_version = toml::toml! {
+            [tool.poetry]
+            name = "starship"
+            version = "0.1.0"
+        }
+        .to_string();
+
+        let expected_version = Some("v0.1.0".to_string());
+        assert_eq!(
+            extract_poetry_version(&poetry_with_version),
+            expected_version
+        );
+
+        let poetry_without_version = toml::toml! {
+            [tool.poetry]
+            name = "starship"
+        }
+        .to_string();
+
+        let expected_version = None;
+        assert_eq!(
+            extract_poetry_version(&poetry_without_version),
+            expected_version
+        );
+    }
 }