Add a method to resolve recursion in optional dependencies

olivier-lacroix · olivier-lacroix · commit 1191fdc97a14 · 2025-04-21T16:09:05.000+10:00
diff --git a/src/has_recursion.rs b/src/has_recursion.rs
@@ -0,0 +1,123 @@
+use indexmap::IndexMap;
+use pep508_rs::Requirement;
+use std::ops::Deref;
+use thiserror::Error;
+
+/// A trait that resolves recursions for groups of requirements that can be mapped to `IndexMap<String, Vec<T>>`
+/// where T is a type that can be mapped to either a Requirement or a reference to other groups of requirements.
+pub trait HasRecursion<T>: Deref<Target = IndexMap<String, Vec<T>>>
+where
+    T: RecursionItem,
+{
+    /// Resolve the groups into lists of requirements.
+    ///
+    /// This function will recursively resolve all groups, including those that
+    /// reference other groups. It will return an error if there is a cycle in the
+    /// groups or if a group references another group that does not exist.
+    fn resolve(&self) -> Result<IndexMap<String, Vec<Requirement>>, RecursionResolutionError> {
+        self.resolve_all(None)
+    }
+
+    /// Resolves the groups of requirements into flat lists of requirements.
+    fn resolve_all(
+        &self,
+        name: Option<&str>,
+    ) -> Result<IndexMap<String, Vec<Requirement>>, RecursionResolutionError> {
+        // Helper function to resolve a single group
+        fn resolve_single<'a, T: RecursionItem>(
+            groups: &'a IndexMap<String, Vec<T>>,
+            group: &'a str,
+            resolved: &mut IndexMap<String, Vec<Requirement>>,
+            parents: &mut Vec<&'a str>,
+            name: Option<&'a str>,
+        ) -> Result<(), RecursionResolutionError> {
+            let Some(items) = groups.get(group) else {
+                // If the group included in another group does not exist, return an error
+                let parent = parents.iter().last().expect("should have a parent");
+                return Err(RecursionResolutionError::GroupNotFound(
+                    T::group_name(),
+                    group.to_string(),
+                    parent.to_string(),
+                ));
+            };
+            // If there is a cycle in dependency groups, return an error
+            if parents.contains(&group) {
+                return Err(RecursionResolutionError::DependencyGroupCycle(
+                    T::table_name(),
+                    Cycle(parents.iter().map(|s| s.to_string()).collect()),
+                ));
+            }
+            // If the group has already been resolved, exit early
+            if resolved.get(group).is_some() {
+                return Ok(());
+            }
+            // Otherwise, perform recursion, as required, on the dependency group's specifiers
+            parents.push(group);
+            let mut requirements = Vec::with_capacity(items.len());
+            for spec in items.iter() {
+                match spec.parse(name) {
+                    // It's a requirement. Just add it to the Vec of resolved requirements
+                    Item::Requirement(requirement) => requirements.push(requirement.clone()),
+                    // It's a reference to other groups. Recurse into them
+                    Item::Groups(inner_groups) => {
+                        for group in inner_groups {
+                            resolve_single(groups, group, resolved, parents, name)?;
+                            requirements.extend(resolved.get(group).into_iter().flatten().cloned());
+                        }
+                    }
+                }
+            }
+            // Add the resolved group to IndexMap
+            resolved.insert(group.to_string(), requirements.clone());
+            parents.pop();
+            Ok(())
+        }
+
+        let mut resolved = IndexMap::new();
+        for group in self.keys() {
+            resolve_single(self, group, &mut resolved, &mut Vec::new(), name)?;
+        }
+        Ok(resolved)
+    }
+}
+/// A trait that defines how to parse a recursion item.
+pub trait RecursionItem {
+    /// Parse the item into a requirement or a reference to other groups.
+    fn parse<'a>(&'a self, name: Option<&str>) -> Item<'a>;
+    /// The name of the group in the TOML file.
+    fn group_name() -> String;
+    /// The name of the table in the TOML file.
+    fn table_name() -> String;
+}
+
+pub enum Item<'a> {
+    Requirement(Requirement),
+    Groups(Vec<&'a str>),
+}
+
+#[derive(Debug, Error)]
+pub enum RecursionResolutionError {
+    #[error("Failed to find {0} `{1}` included by `{2}`")]
+    GroupNotFound(String, String, String),
+    #[error("Detected a cycle in `{0}`: {1}")]
+    DependencyGroupCycle(String, Cycle),
+}
+
+/// A cycle in the recursion.
+#[derive(Debug)]
+pub struct Cycle(Vec<String>);
+
+/// Display a cycle, e.g., `a -> b -> c -> a`.
+impl std::fmt::Display for Cycle {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let [first, rest @ ..] = self.0.as_slice() else {
+            return Ok(());
+        };
+        write!(f, "`{first}`")?;
+        for group in rest {
+            write!(f, " -> `{group}`")?;
+        }
+        write!(f, " -> `{first}`")?;
+        Ok(())
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -4,6 +4,8 @@ mod pep639_glob;
 #[cfg(feature = "pep639-glob")]
 pub use pep639_glob::{parse_pep639_glob, Pep639GlobError};
 
+pub mod has_recursion;
+pub mod optional_dependencies_resolve;
 pub mod pep735_resolve;
 
 use indexmap::IndexMap;
@@ -83,7 +85,7 @@ pub struct Project {
     /// Project dependencies
     pub dependencies: Option<Vec<Requirement>>,
     /// Optional dependencies
-    pub optional_dependencies: Option<IndexMap<String, Vec<Requirement>>>,
+    pub optional_dependencies: Option<OptionalDependencies>,
     /// Specifies which fields listed by PEP 621 were intentionally unspecified
     /// so another tool can/will provide such metadata dynamically.
     pub dynamic: Option<Vec<String>>,
@@ -115,6 +117,23 @@ impl Project {
     }
 }
 
+/// The `[project.optional-dependencies]` section of pyproject.toml
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
+#[serde(transparent)]
+pub struct OptionalDependencies {
+    pub inner: IndexMap<String, Vec<Requirement>>,
+    #[serde(skip)]
+    pub self_reference_name: Option<String>,
+}
+
+impl Deref for OptionalDependencies {
+    type Target = IndexMap<String, Vec<Requirement>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
 /// The full description of the project (i.e. the README).
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
 #[serde(rename_all = "kebab-case")]
@@ -228,7 +247,17 @@ pub enum DependencyGroupSpecifier {
 impl PyProjectToml {
     /// Parse `pyproject.toml` content
     pub fn new(content: &str) -> Result<Self, toml::de::Error> {
-        toml::de::from_str(content)
+        let mut pyproject: PyProjectToml = toml::de::from_str(content)?;
+
+        // Set the project name as optional-dependencies self_reference_name
+        if let Some(project) = pyproject.project.as_mut() {
+            let name = project.name.clone();
+            if let Some(od) = project.optional_dependencies.as_mut() {
+                od.self_reference_name = Some(name);
+            }
+        }
+
+        Ok(pyproject)
     }
 }
 
@@ -337,6 +366,14 @@ tomatoes = "spam:main_tomatoes""#;
             project.gui_scripts.as_ref().unwrap()["spam-gui"],
             "spam:main_gui"
         );
+        assert_eq!(
+            project
+                .optional_dependencies
+                .as_ref()
+                .unwrap()
+                .self_reference_name,
+            Some("spam".to_string())
+        );
     }
 
     #[test]
diff --git a/src/optional_dependencies_resolve.rs b/src/optional_dependencies_resolve.rs
@@ -0,0 +1,111 @@
+use crate::{
+    has_recursion::{HasRecursion, Item, RecursionItem, RecursionResolutionError},
+    OptionalDependencies,
+};
+use indexmap::IndexMap;
+use pep508_rs::Requirement;
+
+impl HasRecursion<Requirement> for OptionalDependencies {
+    fn resolve(&self) -> Result<IndexMap<String, Vec<Requirement>>, RecursionResolutionError> {
+        self.resolve_all(self.self_reference_name.as_deref())
+    }
+}
+
+impl RecursionItem for Requirement {
+    fn parse(&self, name: Option<&str>) -> Item {
+        if name.map(|n| n == self.name.to_string()).unwrap_or(false) {
+            Item::Groups(self.extras.iter().map(|extra| extra.as_ref()).collect())
+        } else {
+            Item::Requirement(self.clone())
+        }
+    }
+    fn table_name() -> String {
+        "project.optional-dependencies".to_string()
+    }
+    fn group_name() -> String {
+        "optional dependency group".to_string()
+    }
+}
+#[cfg(test)]
+mod tests {
+    use pep508_rs::Requirement;
+    use std::str::FromStr;
+
+    use crate::PyProjectToml;
+
+    #[test]
+    fn test_parse_pyproject_toml_dependency_groups_resolve() {
+        let source = r#"[project]
+name = "spam"
+
+[project.optional-dependencies]
+alpha = ["beta", "gamma", "delta"]
+epsilon = ["eta<2.0", "theta==2024.09.01"]
+iota = ["spam[alpha]"]
+"#;
+        let project_toml = PyProjectToml::new(source).unwrap();
+        let optional_dependencies = project_toml
+            .project
+            .as_ref()
+            .unwrap()
+            .optional_dependencies
+            .as_ref()
+            .unwrap();
+
+        assert_eq!(
+            optional_dependencies.resolve().unwrap()["iota"],
+            vec![
+                Requirement::from_str("beta").unwrap(),
+                Requirement::from_str("gamma").unwrap(),
+                Requirement::from_str("delta").unwrap()
+            ]
+        );
+    }
+
+    #[test]
+    fn test_parse_pyproject_toml_dependency_groups_cycle() {
+        let source = r#"[project]
+name = "spam"
+
+[project.optional-dependencies]
+alpha = ["spam[iota]"]
+iota = ["spam[alpha]"]
+"#;
+        let project_toml = PyProjectToml::new(source).unwrap();
+        let optional_dependencies = project_toml
+            .project
+            .as_ref()
+            .unwrap()
+            .optional_dependencies
+            .as_ref()
+            .unwrap();
+        assert_eq!(
+            optional_dependencies.resolve().unwrap_err().to_string(),
+            String::from(
+                "Detected a cycle in `project.optional-dependencies`: `alpha` -> `iota` -> `alpha`"
+            )
+        )
+    }
+
+    #[test]
+    fn test_parse_pyproject_toml_dependency_groups_missing_include() {
+        let source = r#"[project]
+name = "spam"
+
+[project.optional-dependencies]
+iota = ["spam[alpha]"]
+"#;
+        let project_toml = PyProjectToml::new(source).unwrap();
+        let optional_dependencies = project_toml
+            .project
+            .as_ref()
+            .unwrap()
+            .optional_dependencies
+            .as_ref()
+            .unwrap();
+        assert_eq!(
+            optional_dependencies.resolve().unwrap_err().to_string(),
+            String::from("Failed to find optional dependency group `alpha` included by `iota`")
+        )
+    }
+}
diff --git a/src/pep735_resolve.rs b/src/pep735_resolve.rs
