feat: Display Python package version from poetry (#153)

1 files changed, 46 insertions, 8 deletions

diff --git a/src/modules/package.rs b/src/modules/package.rs
index 044499f51..5edacd21e 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
+        );
+    }
 }