Add Distribution::python_version() method

Author: messense

src/distribution.rs

@@ -31,6 +31,7 @@ enum SDistType {
 pub struct Distribution {
     dist_type: DistributionType,
     metadata: Metadata,
+    python_version: String,
 }
 
 impl FromStr for SDistType {
@@ -58,17 +59,44 @@ impl Distribution {
                 "whl" => DistributionType::Wheel,
                 _ => return Err(Error::UnknownDistributionType),
             };
-            let metadata = match dist_type {
+            let (metadata, python_version) = match dist_type {
                 DistributionType::SDist => {
                     let sdist_type: SDistType = ext.parse()?;
-                    Self::parse_sdist(path, sdist_type)
+                    (Self::parse_sdist(path, sdist_type)?, "source".to_string())
                 }
-                DistributionType::Egg => Self::parse_egg(path),
-                DistributionType::Wheel => Self::parse_wheel(path),
-            }?;
+                DistributionType::Egg => {
+                    let parts: Vec<&str> = path
+                        .file_stem()
+                        .unwrap()
+                        .to_str()
+                        .unwrap()
+                        .split('-')
+                        .collect();
+                    let python_version = match parts.as_slice() {
+                        [_name, _version, py_ver] => py_ver,
+                        _ => "any",
+                    };
+                    (Self::parse_egg(path)?, python_version.to_string())
+                }
+                DistributionType::Wheel => {
+                    let parts: Vec<&str> = path
+                        .file_stem()
+                        .unwrap()
+                        .to_str()
+                        .unwrap()
+                        .split('-')
+                        .collect();
+                    let python_version = match parts.as_slice() {
+                        [_name, _version, py_ver, _abi_tag, _plat_tag] => py_ver,
+                        _ => "any",
+                    };
+                    (Self::parse_wheel(path)?, python_version.to_string())
+                }
+            };
             return Ok(Self {
                 dist_type,
                 metadata,
+                python_version,
             });
         }
         Err(Error::UnknownDistributionType)
@@ -84,6 +112,13 @@ impl Distribution {
         &self.metadata
     }
 
+    /// Returns the supported Python version tag
+    ///
+    /// For source distributions the version tag is always `source`
+    pub fn python_version(&self) -> &str {
+        &self.python_version
+    }
+
     fn parse_sdist(path: &Path, sdist_type: SDistType) -> Result<Metadata, Error> {
         match sdist_type {
             SDistType::Zip => Self::parse_zip(path, "PKG-INFO"),

tests/test_distribution.rs

@@ -9,6 +9,7 @@ fn test_parse_wheel() {
     assert_eq!(metadata.name, "build");
     assert!(metadata.home_page.is_none());
     assert!(metadata.download_url.is_none());
+    assert_eq!(dist.python_version(), "py2.py3");
 }
 
 #[test]
@@ -20,6 +21,7 @@ fn test_parse_egg() {
     assert_eq!(metadata.name, "build");
     assert!(metadata.home_page.is_none());
     assert!(metadata.download_url.is_none());
+    assert_eq!(dist.python_version(), "py3.9");
 }
 
 #[test]
@@ -31,6 +33,7 @@ fn test_parse_sdist_zip() {
     assert_eq!(metadata.name, "build");
     assert!(metadata.home_page.is_none());
     assert!(metadata.download_url.is_none());
+    assert_eq!(dist.python_version(), "source");
 }
 
 #[test]
@@ -42,6 +45,7 @@ fn test_parse_sdist_tar_gz() {
     assert_eq!(metadata.name, "build");
     assert!(metadata.home_page.is_none());
     assert!(metadata.download_url.is_none());
+    assert_eq!(dist.python_version(), "source");
 }
 
 #[test]
@@ -53,4 +57,5 @@ fn test_parse_sdist_tar_bz2() {
     assert_eq!(metadata.name, "build");
     assert!(metadata.home_page.is_none());
     assert!(metadata.download_url.is_none());
+    assert_eq!(dist.python_version(), "source");
 }