refactor: Use generics over PyPiDependencies and CondaDependencies (p…

…refix-dev#1303)

src/project/dependencies.rs

@@ -1,79 +1,90 @@
 use indexmap::{Equivalent, IndexMap, IndexSet};
+use itertools::Either;
 use rattler_conda_types::{MatchSpec, NamelessMatchSpec, PackageName};
-use std::{hash::Hash, iter::once};
+use std::{borrow::Cow, hash::Hash};
+
+use super::manifest::{python::PyPiPackageName, PyPiRequirement};
+
+pub type PyPiDependencies = Dependencies<PyPiPackageName, PyPiRequirement>;
+pub type CondaDependencies = Dependencies<PackageName, NamelessMatchSpec>;
 
 /// Holds a list of dependencies where for each package name there can be multiple requirements.
 ///
 /// This is used when combining the dependencies of multiple features. Although each target can only
 /// have one requirement for a given package, when combining the dependencies of multiple features
 /// there can be multiple requirements for a given package.
-#[derive(Default, Debug, Clone)]
-pub struct Dependencies {
-    map: IndexMap<PackageName, IndexSet<NamelessMatchSpec>>,
+///
+/// The generic 'Dependencies' struct is aliased as specific PyPiDependencies and CondaDependencies struct to represent
+/// Pypi and Conda dependencies respectively.
+
+#[derive(Debug, Clone)]
+pub struct Dependencies<N: Hash + Eq + Clone, D: Hash + Eq + Clone> {
+    map: IndexMap<N, IndexSet<D>>,
 }
 
-impl From<IndexMap<PackageName, IndexSet<NamelessMatchSpec>>> for Dependencies {
-    fn from(map: IndexMap<PackageName, IndexSet<NamelessMatchSpec>>) -> Self {
-        Self { map }
+impl<N: Hash + Eq + Clone, D: Hash + Eq + Clone> Default for Dependencies<N, D> {
+    fn default() -> Self {
+        Dependencies {
+            map: IndexMap::new(),
+        }
     }
 }
 
-impl From<IndexMap<PackageName, NamelessMatchSpec>> for Dependencies {
-    fn from(map: IndexMap<PackageName, NamelessMatchSpec>) -> Self {
-        Self {
-            map: map
-                .into_iter()
-                .map(|(k, v)| (k, once(v).collect()))
-                .collect(),
-        }
+impl<N: Hash + Eq + Clone, D: Hash + Eq + Clone> IntoIterator for Dependencies<N, D> {
+    type Item = (N, IndexSet<D>);
+    type IntoIter = indexmap::map::IntoIter<N, IndexSet<D>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.map.into_iter()
     }
 }
 
-impl Dependencies {
-    /// Adds the given spec to the list of dependencies.
-    pub fn insert(&mut self, name: PackageName, spec: NamelessMatchSpec) {
-        self.map.entry(name).or_default().insert(spec);
+impl<'a, M, N: Hash + Eq + Clone + 'a, D: Hash + Eq + Clone + 'a> From<M> for Dependencies<N, D>
+where
+    M: IntoIterator<Item = Cow<'a, IndexMap<N, D>>>,
+{
+    /// Create Dependencies<N, D> from an iterator over items of type Cow<'a, IndexMap<N, D>
+    fn from(m: M) -> Self {
+        m.into_iter().fold(Self::default(), |mut acc: Self, deps| {
+            // Either clone the values from the Cow or move the values from the owned map.
+            let deps_iter = match deps {
+                Cow::Borrowed(borrowed) => Either::Left(
+                    borrowed
+                        .into_iter()
+                        .map(|(name, spec)| (name.clone(), spec.clone())),
+                ),
+                Cow::Owned(owned) => Either::Right(owned.into_iter()),
+            };
+
+            // Add the requirements to the accumulator.
+            for (name, spec) in deps_iter {
+                acc.insert(name, spec);
+            }
+
+            acc
+        })
     }
+}
 
-    /// Adds a list of specs to the list of dependencies.
-    pub fn extend(&mut self, iter: impl IntoIterator<Item = (PackageName, NamelessMatchSpec)>) {
-        for (name, spec) in iter {
-            self.insert(name, spec);
-        }
+impl<N: Hash + Eq + Clone, D: Hash + Eq + Clone> Dependencies<N, D> {
+    /// Adds a requirement to the list of dependencies.
+    pub fn insert(&mut self, name: N, spec: D) {
+        self.map.entry(name).or_default().insert(spec);
     }
 
-    /// Adds a list of specs to the list of dependencies overwriting any existing requirements for
-    /// packages that already exist in the list of dependencies.
-    pub fn extend_overwrite(
-        &mut self,
-        iter: impl IntoIterator<Item = (PackageName, NamelessMatchSpec)>,
-    ) {
-        for (name, spec) in iter {
-            *self.map.entry(name).or_default() = once(spec).collect();
-        }
+    /// Check if there is any dependency
+    pub fn is_empty(&self) -> bool {
+        self.map.is_empty()
     }
 
-    /// Removes all requirements for the given package and returns them.
-    pub fn remove<Q: ?Sized>(
-        &mut self,
-        name: &Q,
-    ) -> Option<(PackageName, IndexSet<NamelessMatchSpec>)>
+    /// Removes a specific dependency
+    pub fn remove<Q: ?Sized>(&mut self, name: &Q) -> Option<(N, IndexSet<D>)>
     where
-        Q: Hash + Equivalent<PackageName>,
+        Q: Hash + Equivalent<N>,
     {
         self.map.shift_remove_entry(name)
     }
 
-    /// Combine two sets of dependencies together where the requirements of `self` are extended if
-    /// the same package is also defined in `other`.
-    pub fn union(&self, other: &Self) -> Self {
-        let mut map = self.map.clone();
-        for (name, specs) in &other.map {
-            map.entry(name.clone()).or_default().extend(specs.clone())
-        }
-        Self { map }
-    }
-
     /// Combines two sets of dependencies where the requirements of `self` are overwritten if the
     /// same package is also defined in `other`.
     pub fn overwrite(&self, other: &Self) -> Self {
@@ -84,34 +95,32 @@ impl Dependencies {
         Self { map }
     }
 
-    /// Returns an iterator over the package names and their corresponding requirements.
-    pub fn iter(
-        &self,
-    ) -> impl DoubleEndedIterator<Item = (&PackageName, &IndexSet<NamelessMatchSpec>)> + '_ {
+    /// Returns an iterator over tuples of dependency names and their combined requirements.
+    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (&N, &IndexSet<D>)> + '_ {
         self.map.iter()
     }
 
-    /// Returns an iterator over all the requirements.
-    pub fn iter_specs(
-        &self,
-    ) -> impl DoubleEndedIterator<Item = (&PackageName, &NamelessMatchSpec)> + '_ {
+    /// Returns an iterator over tuples of dependency names and individual requirements.
+    pub fn iter_specs(&self) -> impl DoubleEndedIterator<Item = (&N, &D)> + '_ {
         self.map
             .iter()
             .flat_map(|(name, specs)| specs.iter().map(move |spec| (name, spec)))
     }
 
-    /// Returns the names of all the packages that have requirements.
-    pub fn names(&self) -> impl DoubleEndedIterator<Item = &PackageName> + ExactSizeIterator + '_ {
+    /// Return an iterator over the dependency names.
+    pub fn names(&self) -> impl DoubleEndedIterator<Item = &N> + ExactSizeIterator + '_ {
         self.map.keys()
     }
 
-    /// Convert this instance into an iterator over the package names and their corresponding
-    pub fn into_specs(self) -> impl DoubleEndedIterator<Item = (PackageName, NamelessMatchSpec)> {
+    /// Converts this instance into an iterator over tuples of dependency names and individual requirements.
+    pub fn into_specs(self) -> impl DoubleEndedIterator<Item = (N, D)> {
         self.map
             .into_iter()
             .flat_map(|(name, specs)| specs.into_iter().map(move |spec| (name.clone(), spec)))
     }
+}
 
+impl Dependencies<PackageName, NamelessMatchSpec> {
     /// Converts this instance into an iterator of [`MatchSpec`]s.
     pub fn into_match_specs(self) -> impl DoubleEndedIterator<Item = MatchSpec> {
         self.map.into_iter().flat_map(|(name, specs)| {
@@ -121,12 +130,3 @@ impl Dependencies {
         })
     }
 }
-
-impl IntoIterator for Dependencies {
-    type Item = (PackageName, IndexSet<NamelessMatchSpec>);
-    type IntoIter = indexmap::map::IntoIter<PackageName, IndexSet<NamelessMatchSpec>>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.map.into_iter()
-    }
-}

src/project/environment.rs

@@ -239,7 +239,7 @@ impl<'p> Hash for Environment<'p> {
 
 #[cfg(test)]
 mod tests {
-    use crate::project::Dependencies;
+    use crate::project::CondaDependencies;
 
     use super::*;
     use insta::assert_snapshot;
@@ -362,7 +362,7 @@ mod tests {
         assert_eq!(task.contains(&"foo".into()), true);
     }
 
-    fn format_dependencies(dependencies: Dependencies) -> String {
+    fn format_dependencies(dependencies: CondaDependencies) -> String {
         dependencies
             .into_specs()
             .map(|(name, spec)| format!("{} = {}", name.as_source(), spec))

src/project/has_features.rs

@@ -1,17 +1,13 @@
-use std::{borrow::Cow, collections::HashSet};
+use std::collections::HashSet;
 
-use indexmap::{IndexMap, IndexSet};
-use itertools::Either;
+use indexmap::IndexSet;
 use rattler_conda_types::{Channel, Platform};
 
 use crate::{Project, SpecType};
 
 use super::{
-    manifest::{
-        pypi_options::PypiOptions, python::PyPiPackageName, Feature, PyPiRequirement,
-        SystemRequirements,
-    },
-    Dependencies,
+    manifest::{pypi_options::PypiOptions, Feature, SystemRequirements},
+    CondaDependencies, PyPiDependencies,
 };
 
 /// A trait that implement various methods for collections that combine attributes of Features
@@ -102,43 +98,25 @@ pub trait HasFeatures<'p> {
     /// The dependencies of all features are combined. This means that if two features define a
     /// requirement for the same package that both requirements are returned. The different
     /// requirements per package are sorted in the same order as the features they came from.
-    fn pypi_dependencies(
-        &self,
-        platform: Option<Platform>,
-    ) -> IndexMap<PyPiPackageName, IndexSet<PyPiRequirement>> {
+    fn pypi_dependencies(&self, platform: Option<Platform>) -> PyPiDependencies {
         self.features()
             .filter_map(|f| f.pypi_dependencies(platform))
-            .fold(IndexMap::default(), |mut acc, deps| {
-                // Either clone the values from the Cow or move the values from the owned map.
-                let deps_iter = match deps {
-                    Cow::Borrowed(borrowed) => Either::Left(
-                        borrowed
-                            .into_iter()
-                            .map(|(name, spec)| (name.clone(), spec.clone())),
-                    ),
-                    Cow::Owned(owned) => Either::Right(owned.into_iter()),
-                };
-
-                // Add the requirements to the accumulator.
-                for (name, spec) in deps_iter {
-                    acc.entry(name).or_default().insert(spec);
-                }
-
-                acc
-            })
+            .into()
     }
 
     /// Returns the dependencies to install for this collection.
     ///
     /// The dependencies of all features are combined. This means that if two features define a
     /// requirement for the same package that both requirements are returned. The different
     /// requirements per package are sorted in the same order as the features they came from.
-    fn dependencies(&self, kind: Option<SpecType>, platform: Option<Platform>) -> Dependencies {
+    fn dependencies(
+        &self,
+        kind: Option<SpecType>,
+        platform: Option<Platform>,
+    ) -> CondaDependencies {
         self.features()
             .filter_map(|f| f.dependencies(kind, platform))
-            .map(|deps| Dependencies::from(deps.into_owned()))
-            .reduce(|acc, deps| acc.union(&deps))
-            .unwrap_or_default()
+            .into()
     }
 
     /// Returns the pypi options for this solve group.

src/project/mod.rs

@@ -8,8 +8,7 @@ mod solve_group;
 pub mod virtual_packages;
 
 use async_once_cell::OnceCell as AsyncCell;
-
-use indexmap::{Equivalent, IndexMap, IndexSet};
+use indexmap::{Equivalent, IndexSet};
 use miette::{IntoDiagnostic, NamedSource};
 
 use rattler_conda_types::{Channel, Platform, Version};
@@ -35,17 +34,15 @@ use crate::{
     consts::{self, PROJECT_MANIFEST, PYPROJECT_MANIFEST},
     task::Task,
 };
-use manifest::{EnvironmentName, Manifest, PyPiRequirement, SystemRequirements};
-
-use crate::project::manifest::python::PyPiPackageName;
-pub use dependencies::Dependencies;
-pub use environment::Environment;
-pub use solve_group::SolveGroup;
+use manifest::{EnvironmentName, Manifest, SystemRequirements};
 
 use self::{
     has_features::HasFeatures,
     manifest::{pyproject::PyProjectToml, Environments},
 };
+pub use dependencies::{CondaDependencies, PyPiDependencies};
+pub use environment::Environment;
+pub use solve_group::SolveGroup;
 
 /// The dependency types we support
 #[derive(Debug, Copy, Clone)]
@@ -437,17 +434,18 @@ impl Project {
     /// Returns the dependencies of the project.
     ///
     /// TODO: Remove this function and use the `dependencies` function from the default environment instead.
-    pub fn dependencies(&self, kind: Option<SpecType>, platform: Option<Platform>) -> Dependencies {
+    pub fn dependencies(
+        &self,
+        kind: Option<SpecType>,
+        platform: Option<Platform>,
+    ) -> CondaDependencies {
         self.default_environment().dependencies(kind, platform)
     }
 
     /// Returns the PyPi dependencies of the project
     ///
     /// TODO: Remove this function and use the `dependencies` function from the default environment instead.
-    pub fn pypi_dependencies(
-        &self,
-        platform: Option<Platform>,
-    ) -> IndexMap<PyPiPackageName, IndexSet<PyPiRequirement>> {
+    pub fn pypi_dependencies(&self, platform: Option<Platform>) -> PyPiDependencies {
         self.default_environment().pypi_dependencies(platform)
     }
 
@@ -597,7 +595,7 @@ mod tests {
         }
     }
 
-    fn format_dependencies(deps: Dependencies) -> String {
+    fn format_dependencies(deps: CondaDependencies) -> String {
         deps.iter_specs()
             .map(|(name, spec)| format!("{} = \"{}\"", name.as_source(), spec))
             .join("\n")