uv_resolver/

flat_index.rs

use std::collections::BTreeMap;
use std::collections::btree_map::Entry;

use rustc_hash::FxHashMap;
use tracing::instrument;

use uv_cache::Cache;
use uv_client::{
    FlatIndexClient, FlatIndexEntries, FlatIndexEntry, FlatIndexError, RegistryClient,
};
use uv_configuration::BuildOptions;
use uv_distribution_filename::{DistFilename, SourceDistFilename, WheelFilename};
use uv_distribution_types::{
    File, HashComparison, IncompatibleSource, IncompatibleWheel, Index, IndexLocations, IndexUrl,
    PrioritizedDist, RegistryBuiltWheel, RegistrySourceDist, SourceDistCompatibility,
    WheelCompatibility,
};
use uv_normalize::PackageName;
use uv_pep440::Version;
use uv_platform_tags::{TagCompatibility, Tags};
use uv_pypi_types::HashDigest;
use uv_types::HashStrategy;

/// Unfiltered entries from `--find-links`, indexed by [`PackageName`].
#[derive(Debug, Clone, Default)]
pub struct FlatIndex {
    /// Entries are ranked using each resolver's platform and hash policies when requested.
    index: FxHashMap<PackageName, Vec<FlatIndexEntry>>,
    /// Whether any `--find-links` entries could not be resolved due to a lack of network
    /// connectivity.
    offline: bool,
}

impl FlatIndex {
    /// Load the `--find-links` entries from the configured indexes.
    pub async fn load(
        client: &RegistryClient,
        cache: &Cache,
        index_locations: &IndexLocations,
    ) -> Result<Self, FlatIndexError> {
        let client = FlatIndexClient::new(client.cached_client(), client.connectivity(), cache);
        let entries = client
            .fetch_all(index_locations.flat_indexes().map(Index::url))
            .await?;
        Ok(Self::from_entries(entries))
    }

    /// Collect all files from a `--find-links` target into a [`FlatIndex`].
    #[instrument(skip_all)]
    fn from_entries(entries: FlatIndexEntries) -> Self {
        let mut index = FxHashMap::<PackageName, Vec<FlatIndexEntry>>::default();
        let (entries, offline) = entries.into_parts();

        for entry in entries {
            index
                .entry(entry.filename().name().clone())
                .or_default()
                .push(entry);
        }

        Self { index, offline }
    }

    /// Get the unfiltered entries for the given package name.
    pub(crate) fn get(&self, package_name: &PackageName) -> Option<&[FlatIndexEntry]> {
        self.index.get(package_name).map(Vec::as_slice)
    }

    /// Whether any `--find-links` entries could not be resolved due to a lack of network
    /// connectivity.
    pub(crate) fn offline(&self) -> bool {
        self.offline
    }
}

/// A set of [`PrioritizedDist`] from a `--find-links` entry for a single package, indexed
/// by [`Version`].
#[derive(Debug, Clone, Default)]
pub(crate) struct FlatDistributions(BTreeMap<Version, PrioritizedDist>);

impl FlatDistributions {
    /// Rank the entries for a package using the resolver's platform and hash policies.
    #[instrument(skip_all)]
    pub(crate) fn from_entries(
        entries: impl IntoIterator<Item = FlatIndexEntry>,
        tags: Option<&Tags>,
        hasher: &HashStrategy,
        build_options: &BuildOptions,
    ) -> Self {
        let mut distributions = Self::default();
        for entry in entries {
            let (filename, file, index) = entry.into_parts();
            distributions.add_file(file, filename, tags, hasher, build_options, index);
        }
        distributions
    }

    /// Returns an [`Iterator`] over the distributions.
    pub(crate) fn iter(&self) -> impl Iterator<Item = (&Version, &PrioritizedDist)> {
        self.0.iter()
    }

    /// Add the given [`File`] to the [`FlatDistributions`] for the given package.
    fn add_file(
        &mut self,
        file: File,
        filename: DistFilename,
        tags: Option<&Tags>,
        hasher: &HashStrategy,
        build_options: &BuildOptions,
        index: IndexUrl,
    ) {
        // No `requires-python` here: for source distributions, we don't have that information;
        // for wheels, we read it lazily only when selected.
        match filename {
            DistFilename::WheelFilename(filename) => {
                let version = filename.version.clone();

                let compatibility = Self::wheel_compatibility(
                    &filename,
                    file.hashes.as_slice(),
                    tags,
                    hasher,
                    build_options,
                );
                let dist = RegistryBuiltWheel {
                    filename,
                    file: Box::new(file),
                    index,
                    size_is_authoritative: false,
                };
                match self.0.entry(version) {
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert_built(dist, vec![], compatibility);
                    }
                    Entry::Vacant(entry) => {
                        entry.insert(PrioritizedDist::from_built(dist, vec![], compatibility));
                    }
                }
            }
            DistFilename::SourceDistFilename(filename) => {
                let compatibility = Self::source_dist_compatibility(
                    &filename,
                    file.hashes.as_slice(),
                    hasher,
                    build_options,
                );
                let dist = RegistrySourceDist {
                    name: filename.name.clone(),
                    version: filename.version.clone(),
                    ext: filename.extension,
                    file: Box::new(file),
                    index,
                    wheels: vec![],
                    size_is_authoritative: false,
                };
                match self.0.entry(filename.version) {
                    Entry::Occupied(mut entry) => {
                        entry.get_mut().insert_source(dist, vec![], compatibility);
                    }
                    Entry::Vacant(entry) => {
                        entry.insert(PrioritizedDist::from_source(dist, vec![], compatibility));
                    }
                }
            }
        }
    }

    fn source_dist_compatibility(
        filename: &SourceDistFilename,
        hashes: &[HashDigest],
        hasher: &HashStrategy,
        build_options: &BuildOptions,
    ) -> SourceDistCompatibility {
        // Check if source distributions are allowed for this package.
        if build_options.no_build_package(&filename.name) {
            return SourceDistCompatibility::Incompatible(IncompatibleSource::NoBuild);
        }

        // Check if the filename is PEP 625-compliant.
        // TODO: Strengthen this check more; right now we allow `.zip`
        // (which is not compliant) and we don't strictly
        // enforce the formatting rules for the name or version.
        if !filename.extension.is_pep625_compliant() {
            return SourceDistCompatibility::Incompatible(IncompatibleSource::NotPep625Filename);
        }

        // Check if hashes line up
        let hash_policy = hasher.archive_policy_for_package(&filename.name, &filename.version);
        let hash = if hash_policy.requires_validation() {
            if hashes.is_empty() {
                HashComparison::Missing
            } else if hash_policy.matches(hashes) {
                HashComparison::Matched
            } else {
                HashComparison::Mismatched
            }
        } else {
            HashComparison::Matched
        };

        SourceDistCompatibility::Compatible(hash)
    }

    fn wheel_compatibility(
        filename: &WheelFilename,
        hashes: &[HashDigest],
        tags: Option<&Tags>,
        hasher: &HashStrategy,
        build_options: &BuildOptions,
    ) -> WheelCompatibility {
        // Check if binaries are allowed for this package.
        if build_options.no_binary_package(&filename.name) {
            return WheelCompatibility::Incompatible(IncompatibleWheel::NoBinary);
        }

        // Determine a compatibility for the wheel based on tags.
        let priority = match tags {
            Some(tags) => match filename.compatibility(tags) {
                TagCompatibility::Incompatible(tag) => {
                    return WheelCompatibility::Incompatible(IncompatibleWheel::Tag(tag));
                }
                TagCompatibility::Compatible(priority) => Some(priority),
            },
            None => None,
        };

        // Check if hashes line up.
        let hash_policy = hasher.archive_policy_for_package(&filename.name, &filename.version);
        let hash = if hash_policy.requires_validation() {
            if hashes.is_empty() {
                HashComparison::Missing
            } else if hash_policy.matches(hashes) {
                HashComparison::Matched
            } else {
                HashComparison::Mismatched
            }
        } else {
            HashComparison::Matched
        };

        // Break ties with the build tag.
        let build_tag = filename.build_tag().cloned();

        WheelCompatibility::Compatible(hash, priority, build_tag)
    }
}

impl IntoIterator for FlatDistributions {
    type Item = (Version, PrioritizedDist);
    type IntoIter = std::collections::btree_map::IntoIter<Version, PrioritizedDist>;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl From<FlatDistributions> for BTreeMap<Version, PrioritizedDist> {
    fn from(distributions: FlatDistributions) -> Self {
        distributions.0
    }
}