refactor: extract pypi-dependencies from environment (#656)

Refactors the code to get the `pypi-dependencies` from `Environment`
instead of directly from the manifest.

---------

Co-authored-by: Ruben Arts <[email protected]>

Cargo.toml

@@ -74,6 +74,7 @@ signal-hook = "0.3.17"
 
 [dev-dependencies]
 rattler_digest = "0.16.2"
+rstest = "0.18.2"
 serde_json = "1.0.111"
 serial_test = "2.0.0"
 tokio = { version = "1.35.1", features = ["rt"] }

src/lock_file/pypi.rs

@@ -22,7 +22,7 @@ pub async fn resolve_dependencies<'p>(
     platform: Platform,
     conda_packages: &[RepoDataRecord],
 ) -> miette::Result<Vec<PinnedPackage<'p>>> {
-    let dependencies = project.pypi_dependencies(platform);
+    let dependencies = project.pypi_dependencies(Some(platform));
     if dependencies.is_empty() {
         return Ok(vec![]);
     }
@@ -70,6 +70,7 @@ pub async fn resolve_dependencies<'p>(
 
     let requirements = dependencies
         .iter()
+        .flat_map(|(name, req)| req.iter().map(move |req| (name, req)))
         .map(|(name, req)| req.as_pep508(name))
         .collect::<Vec<pep508_rs::Requirement>>();
 

src/lock_file/satisfiability.rs

@@ -53,9 +53,11 @@ pub fn lock_file_satisfies_project(
             .collect::<Vec<_>>();
 
         let mut pypi_dependencies = project
-            .pypi_dependencies(platform)
+            .pypi_dependencies(Some(platform))
             .into_iter()
-            .map(|(name, requirement)| requirement.as_pep508(&name))
+            .flat_map(|(name, requirement)| {
+                requirement.into_iter().map(move |req| req.as_pep508(&name))
+            })
             .map(DependencyKind::PyPi)
             .peekable();
 

src/project/environment.rs

@@ -2,12 +2,14 @@ use super::{
     dependencies::Dependencies,
     errors::{UnknownTask, UnsupportedPlatformError},
     manifest::{self, EnvironmentName, Feature, FeatureName, SystemRequirements},
-    SpecType,
+    PyPiRequirement, SpecType,
 };
 use crate::{task::Task, Project};
-use indexmap::IndexSet;
+use indexmap::{IndexMap, IndexSet};
+use itertools::Either;
 use rattler_conda_types::{Channel, Platform};
 use std::{
+    borrow::Cow,
     collections::{HashMap, HashSet},
     fmt::Debug,
 };
@@ -174,31 +176,48 @@ impl<'p> Environment<'p> {
 
     /// Returns the dependencies to install for this environment.
     ///
-    /// The dependencies of all features are combined this means that if two features define a
+    /// 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 from the most specific feature/target to the least
-    /// specific.
+    /// requirements per package are sorted in the same order as the features they came from.
     pub fn dependencies(&self, kind: Option<SpecType>, platform: Option<Platform>) -> Dependencies {
         self.features()
-            .filter_map(|f| {
-                f.targets
-                    .resolve(platform)
-                    .rev()
-                    .map(|t| t.dependencies(kind))
-                    .fold(None, |acc: Option<Dependencies>, deps| {
-                        Some(match acc {
-                            None => Dependencies::from(deps.into_owned()),
-                            Some(mut acc) => {
-                                acc.extend_overwrite(deps.into_owned());
-                                acc
-                            }
-                        })
-                    })
-            })
+            .filter_map(|f| f.dependencies(kind, platform))
+            .map(|deps| Dependencies::from(deps.into_owned()))
             .reduce(|acc, deps| acc.union(&deps))
             .unwrap_or_default()
     }
 
+    /// Returns the PyPi dependencies to install for this environment.
+    ///
+    /// 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.
+    pub fn pypi_dependencies(
+        &self,
+        platform: Option<Platform>,
+    ) -> IndexMap<rip::types::PackageName, Vec<PyPiRequirement>> {
+        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().push(spec);
+                }
+
+                acc
+            })
+    }
+
     /// Validates that the given platform is supported by this environment.
     fn validate_platform_support(
         &self,

src/project/manifest/feature.rs

@@ -4,10 +4,12 @@ use crate::project::SpecType;
 use crate::task::Task;
 use crate::utils::spanned::PixiSpanned;
 use indexmap::IndexMap;
+use itertools::Either;
 use rattler_conda_types::{Channel, NamelessMatchSpec, PackageName, Platform};
 use serde::de::Error;
 use serde::{Deserialize, Deserializer};
 use serde_with::{serde_as, DisplayFromStr, PickFirst};
+use std::borrow::Cow;
 use std::collections::HashMap;
 
 /// The name of a feature. This is either a string or default for the default feature.
@@ -70,6 +72,73 @@ pub struct Feature {
     pub targets: Targets,
 }
 
+impl Feature {
+    /// Returns the dependencies of the feature for a given `spec_type` and `platform`.
+    ///
+    /// This function returns a [`Cow`]. If the dependencies are not combined or overwritten by
+    /// multiple targets than this function returns a reference to the internal dependencies.
+    ///
+    /// Returns `None` if this feature does not define any target that has any of the requested
+    /// dependencies.
+    pub fn dependencies(
+        &self,
+        spec_type: Option<SpecType>,
+        platform: Option<Platform>,
+    ) -> Option<Cow<'_, IndexMap<PackageName, NamelessMatchSpec>>> {
+        self.targets
+            .resolve(platform)
+            // Get the targets in reverse order, from least specific to most specific.
+            // This is required because the extend function will overwrite existing keys.
+            .rev()
+            .filter_map(|t| t.dependencies(spec_type))
+            .filter(|deps| !deps.is_empty())
+            .fold(None, |acc, deps| match acc {
+                None => Some(deps),
+                Some(mut acc) => {
+                    let deps_iter = match deps {
+                        Cow::Borrowed(deps) => Either::Left(
+                            deps.iter().map(|(name, spec)| (name.clone(), spec.clone())),
+                        ),
+                        Cow::Owned(deps) => Either::Right(deps.into_iter()),
+                    };
+
+                    acc.to_mut().extend(deps_iter);
+                    Some(acc)
+                }
+            })
+    }
+
+    /// Returns the PyPi dependencies of the feature for a given `platform`.
+    ///
+    /// This function returns a [`Cow`]. If the dependencies are not combined or overwritten by
+    /// multiple targets than this function returns a reference to the internal dependencies.
+    ///
+    /// Returns `None` if this feature does not define any target that has any of the requested
+    /// dependencies.
+    pub fn pypi_dependencies(
+        &self,
+        platform: Option<Platform>,
+    ) -> Option<Cow<'_, IndexMap<rip::types::PackageName, PyPiRequirement>>> {
+        self.targets
+            .resolve(platform)
+            // Get the targets in reverse order, from least specific to most specific.
+            // This is required because the extend function will overwrite existing keys.
+            .rev()
+            .filter_map(|t| t.pypi_dependencies.as_ref())
+            .filter(|deps| !deps.is_empty())
+            .fold(None, |acc, deps| match acc {
+                None => Some(Cow::Borrowed(deps)),
+                Some(mut acc) => {
+                    acc.to_mut().extend(
+                        deps.into_iter()
+                            .map(|(name, spec)| (name.clone(), spec.clone())),
+                    );
+                    Some(acc)
+                }
+            })
+    }
+}
+
 impl<'de> Deserialize<'de> for Feature {
     fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
     where
@@ -138,3 +207,78 @@ impl<'de> Deserialize<'de> for Feature {
         })
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::project::manifest::Manifest;
+    use assert_matches::assert_matches;
+    use std::path::Path;
+
+    #[test]
+    fn test_dependencies_borrowed() {
+        let manifest = Manifest::from_str(
+            Path::new(""),
+            r#"
+        [project]
+        name = "foo"
+        platforms = ["linux-64", "osx-64", "win-64"]
+        channels = []
+
+        [dependencies]
+        foo = "1.0"
+
+        [host-dependencies]
+        foo = "2.0"
+
+        [feature.bla.dependencies]
+        foo = "2.0"
+
+        [feature.bla.host-dependencies]
+        # empty on purpose
+        "#,
+        )
+        .unwrap();
+
+        assert_matches!(
+            manifest
+                .default_feature()
+                .dependencies(Some(SpecType::Host), None)
+                .unwrap(),
+            Cow::Borrowed(_),
+            "[host-dependencies] should be borrowed"
+        );
+
+        assert_matches!(
+            manifest
+                .default_feature()
+                .dependencies(Some(SpecType::Run), None)
+                .unwrap(),
+            Cow::Borrowed(_),
+            "[dependencies] should be borrowed"
+        );
+
+        assert_matches!(
+            manifest.default_feature().dependencies(None, None).unwrap(),
+            Cow::Owned(_),
+            "combined dependencies should be owned"
+        );
+
+        let bla_feature = manifest
+            .parsed
+            .features
+            .get(&FeatureName::Named(String::from("bla")))
+            .unwrap();
+        assert_matches!(
+            bla_feature.dependencies(Some(SpecType::Run), None).unwrap(),
+            Cow::Borrowed(_),
+            "[feature.bla.dependencies] should be borrowed"
+        );
+
+        assert_matches!(
+            bla_feature.dependencies(None, None).unwrap(),
+            Cow::Borrowed(_),
+            "[feature.bla] combined dependencies should also be borrowed"
+        );
+    }
+}