Add index field to pypi requirements (#784)

prefix-dev · Feb 8, 2024 · af8f8cf · af8f8cf
1 parent a9f1f61
commit af8f8cf
Showing 1 changed file with 29 additions and 8 deletions.
diff --git a/src/project/manifest/python.rs b/src/project/manifest/python.rs
@@ -8,6 +8,7 @@ use toml_edit::Item;
 pub struct PyPiRequirement {
     pub(crate) version: Option<pep440_rs::VersionSpecifiers>,
     pub(crate) extras: Option<Vec<String>>,
+    pub(crate) index: Option<String>,
 }
 
 /// The type of parse error that occurred when parsing match spec.
@@ -80,6 +81,7 @@ impl FromStr for PyPiRequirement {
             Ok(Self {
                 version: None,
                 extras: None,
+                index: None,
             })
         } else if s.starts_with(|c: char| c.is_ascii_digit()) {
             Err(ParsePyPiRequirementError::MissingOperator(s.to_string()))
@@ -91,6 +93,7 @@ impl FromStr for PyPiRequirement {
                         .map_err(ParsePyPiRequirementError::Pep440Error)?,
                 ),
                 extras: None,
+                index: None,
             })
         }
     }
@@ -110,6 +113,7 @@ impl From<pep508_rs::Requirement> for PyPiRequirement {
         PyPiRequirement {
             version,
             extras: req.extras,
+            index: None,
         }
     }
 }
@@ -141,6 +145,7 @@ impl<'de> Deserialize<'de> for PyPiRequirement {
                 struct RawPyPiRequirement {
                     version: Option<String>,
                     extras: Option<Vec<String>>,
+                    index: Option<String>,
                 }
                 let raw_requirement =
                     RawPyPiRequirement::deserialize(de::value::MapAccessDeserializer::new(map))?;
@@ -157,6 +162,7 @@ impl<'de> Deserialize<'de> for PyPiRequirement {
                 Ok(PyPiRequirement {
                     version,
                     extras: raw_requirement.extras,
+                    index: raw_requirement.index,
                 })
             })
             .expecting("either a map or a string")
@@ -192,7 +198,8 @@ mod tests {
             requirement.first().unwrap().1,
             &PyPiRequirement {
                 version: Some(pep440_rs::VersionSpecifiers::from_str(">=3.12").unwrap()),
-                extras: None
+                extras: None,
+                index: None,
             }
         );
         let requirement: IndexMap<rip::types::PackageName, PyPiRequirement> =
@@ -201,7 +208,8 @@ mod tests {
             requirement.first().unwrap().1,
             &PyPiRequirement {
                 version: Some(pep440_rs::VersionSpecifiers::from_str("==3.12.0").unwrap()),
-                extras: None
+                extras: None,
+                index: None,
             }
         );
 
@@ -211,7 +219,8 @@ mod tests {
             requirement.first().unwrap().1,
             &PyPiRequirement {
                 version: Some(pep440_rs::VersionSpecifiers::from_str("~=2.1.3").unwrap()),
-                extras: None
+                extras: None,
+                index: None,
             }
         );
 
@@ -221,15 +230,22 @@ mod tests {
             requirement.first().unwrap().1,
             &PyPiRequirement {
                 version: None,
-                extras: None
+                extras: None,
+                index: None,
             }
         );
     }
 
     #[test]
     fn test_extended() {
         let requirement: IndexMap<rip::types::PackageName, PyPiRequirement> =
-            toml_edit::de::from_str(r#"foo = { version=">=3.12", extras = ["bar"] }"#).unwrap();
+            toml_edit::de::from_str(
+                r#"
+                foo = { version=">=3.12", extras = ["bar"], index = "artifact-registry" }
+                "#,
+            )
+            .unwrap();
+
         assert_eq!(
             requirement.first().unwrap().0,
             &rip::types::PackageName::from_str("foo").unwrap()
@@ -238,7 +254,8 @@ mod tests {
             requirement.first().unwrap().1,
             &PyPiRequirement {
                 version: Some(pep440_rs::VersionSpecifiers::from_str(">=3.12").unwrap()),
-                extras: Some(vec!("bar".to_string()))
+                extras: Some(vec!("bar".to_string())),
+                index: Some("artifact-registry".to_string()),
             }
         );
 
@@ -255,7 +272,8 @@ mod tests {
             requirement.first().unwrap().1,
             &PyPiRequirement {
                 version: Some(pep440_rs::VersionSpecifiers::from_str(">=3.12,<3.13.0").unwrap()),
-                extras: Some(vec!("bar".to_string(), "foo".to_string()))
+                extras: Some(vec!("bar".to_string(), "foo".to_string())),
+                index: None,
             }
         );
     }
@@ -275,7 +293,8 @@ mod tests {
             pypi_requirement,
             PyPiRequirement {
                 version: VersionSpecifiers::from_str("==1.2.3").ok(),
-                extras: Some(vec!["feature1".to_owned(), "feature2".to_owned()])
+                extras: Some(vec!["feature1".to_owned(), "feature2".to_owned()]),
+                index: None,
             }
         );
     }
@@ -291,6 +310,7 @@ mod tests {
             PyPiRequirement {
                 version: VersionSpecifiers::from_str("==1.2.3").ok(),
                 extras: None,
+                index: None,
             }
         );
     }
@@ -306,6 +326,7 @@ mod tests {
             PyPiRequirement {
                 version: None,
                 extras: None,
+                index: None,
             }
         );
     }