uv_requirements/
extras.rs

1use std::sync::Arc;
2
3use futures::{TryStreamExt, stream::FuturesOrdered};
4
5use uv_distribution::{DistributionDatabase, Reporter};
6use uv_distribution_types::DistributionMetadata;
7use uv_distribution_types::Requirement;
8use uv_resolver::{InMemoryIndex, MetadataResponse};
9use uv_types::{BuildContext, HashStrategy};
10
11use crate::{Error, required_dist};
12
13/// A resolver to expand the requested extras for a set of requirements to include all defined
14/// extras.
15pub struct ExtrasResolver<'a, Context: BuildContext> {
16    /// Whether to check hashes for distributions.
17    hasher: &'a HashStrategy,
18    /// The in-memory index for resolving dependencies.
19    index: &'a InMemoryIndex,
20    /// The database for fetching and building distributions.
21    database: DistributionDatabase<'a, Context>,
22}
23
24impl<'a, Context: BuildContext> ExtrasResolver<'a, Context> {
25    /// Instantiate a new [`ExtrasResolver`] for a given set of requirements.
26    pub fn new(
27        hasher: &'a HashStrategy,
28        index: &'a InMemoryIndex,
29        database: DistributionDatabase<'a, Context>,
30    ) -> Self {
31        Self {
32            hasher,
33            index,
34            database,
35        }
36    }
37
38    /// Set the [`Reporter`] to use for this resolver.
39    #[must_use]
40    pub fn with_reporter(self, reporter: Arc<dyn Reporter>) -> Self {
41        Self {
42            database: self.database.with_reporter(reporter),
43            ..self
44        }
45    }
46
47    /// Expand the set of available extras for a given set of requirements.
48    pub async fn resolve(
49        self,
50        requirements: impl Iterator<Item = Requirement>,
51    ) -> Result<Vec<Requirement>, Error> {
52        let Self {
53            hasher,
54            index,
55            database,
56        } = self;
57        requirements
58            .map(async |requirement| {
59                Self::resolve_requirement(requirement, hasher, index, &database).await
60            })
61            .collect::<FuturesOrdered<_>>()
62            .try_collect()
63            .await
64    }
65
66    /// Expand the set of available extras for a given [`Requirement`].
67    async fn resolve_requirement(
68        requirement: Requirement,
69        hasher: &HashStrategy,
70        index: &InMemoryIndex,
71        database: &DistributionDatabase<'a, Context>,
72    ) -> Result<Requirement, Error> {
73        // Determine whether the requirement represents a local distribution and convert to a
74        // buildable distribution.
75        let Some(dist) = required_dist(&requirement)? else {
76            return Ok(requirement);
77        };
78
79        // Fetch the metadata for the distribution.
80        let metadata = {
81            let id = dist.version_id();
82            if let Some(archive) = index
83                .distributions()
84                .get(&id)
85                .as_deref()
86                .and_then(|response| {
87                    if let MetadataResponse::Found(archive, ..) = response {
88                        Some(archive)
89                    } else {
90                        None
91                    }
92                })
93            {
94                // If the metadata is already in the index, return it.
95                archive.metadata.clone()
96            } else {
97                // Run the PEP 517 build process to extract metadata from the source distribution.
98                let archive = database
99                    .get_or_build_wheel_metadata(&dist, hasher.get(&dist))
100                    .await
101                    .map_err(|err| Error::from_dist(dist, err))?;
102
103                let metadata = archive.metadata.clone();
104
105                // Insert the metadata into the index.
106                index
107                    .distributions()
108                    .done(id, Arc::new(MetadataResponse::Found(archive)));
109
110                metadata
111            }
112        };
113
114        // Sort extras for consistency.
115        let extras = {
116            let mut extras = metadata.provides_extra.to_vec();
117            extras.sort_unstable();
118            extras
119        };
120
121        Ok(Requirement {
122            extras: extras.into_boxed_slice(),
123            ..requirement
124        })
125    }
126}
uv_requirements/extras.rs

uv_requirements/
extras.rs
