rattler_lock 0.30.3

//! A module that enables parsing of lock files version 3 or lower.

use std::{collections::BTreeSet, ops::Not, str::FromStr, sync::Arc};

use super::{
    ParseCondaLockError,
    models::legacy::{LegacyCondaBinaryData, LegacyCondaPackageData},
};
use crate::{
    Channel, DEFAULT_ENVIRONMENT_NAME, EnvironmentData, EnvironmentIndex, EnvironmentPackages,
    LockFile, LockFileInner, LockedPackage, PackageHashes, PackageIndex, PlatformData,
    PlatformIndex, PypiDistributionData, PypiPackageData, SolveOptions, UrlOrPath, Verbatim,
    file_format_version::FileFormatVersion,
    platform::PlatformName,
    utils::derived_fields::{
        LocationDerivedFields, derive_arch_and_platform, derive_build_number_from_build,
        derive_noarch_type,
    },
};
use indexmap::IndexSet;
use pep440_rs::VersionSpecifiers;
use pep508_rs::Requirement;
use rattler_conda_types::{NoArchType, PackageName, PackageRecord, PackageUrl, VersionWithSource};
use serde::Deserialize;
use serde_with::{OneOrMany, serde_as, skip_serializing_none};
use url::Url;

#[derive(Deserialize)]
struct LockFileV3 {
    metadata: LockMetaV3,
    package: Vec<LockedPackageV3>,
}

#[serde_as]
#[derive(Deserialize, Clone, Debug, Eq, PartialEq)]
struct LockMetaV3 {
    /// Channels used to resolve dependencies
    pub channels: Vec<Channel>,
    /// The platforms this lock file supports
    #[serde_as(as = "crate::utils::serde::Ordered<_>")]
    pub platforms: Vec<rattler_conda_types::Platform>,
}

#[derive(Deserialize, Eq, PartialEq, Clone, Debug)]
struct LockedPackageV3 {
    pub platform: rattler_conda_types::Platform,
    #[serde(flatten)]
    pub kind: LockedPackageKindV3,
}

#[derive(Deserialize, Eq, PartialEq, Clone, Debug)]
#[serde(tag = "manager", rename_all = "snake_case")]
enum LockedPackageKindV3 {
    Conda(Box<CondaLockedPackageV3>),
    #[serde(alias = "pip")]
    Pypi(Box<PypiLockedPackageV3>),
}

#[serde_as]
#[skip_serializing_none]
#[derive(Deserialize, Eq, PartialEq, Clone, Debug)]
struct PypiLockedPackageV3 {
    pub name: String,
    pub version: pep440_rs::Version,
    #[serde(default, alias = "dependencies", skip_serializing_if = "Vec::is_empty")]
    #[serde_as(deserialize_as = "crate::utils::serde::Pep440MapOrVec")]
    pub requires_dist: Vec<Requirement>,
    pub requires_python: Option<VersionSpecifiers>,
    pub url: Url,
    pub hash: Option<PackageHashes>,
    // These fields are not used by rattler-lock.
    // pub source: Option<Url>,
    // pub build: Option<String>,
}

#[serde_as]
#[skip_serializing_none]
#[derive(Deserialize, Eq, PartialEq, Clone, Debug)]
pub struct CondaLockedPackageV3 {
    pub name: String,
    pub version: VersionWithSource,
    #[serde(default)]
    #[serde_as(deserialize_as = "crate::utils::serde::MatchSpecMapOrVec")]
    pub dependencies: Vec<String>,
    pub url: Url,
    pub hash: PackageHashes,
    pub source: Option<Url>,
    pub build: Option<String>,
    pub arch: Option<String>,
    // Platform is used to indicate the actual platform to which this package belongs.
    // pub platform: Option<String>,
    pub subdir: Option<String>,
    pub build_number: Option<u64>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub constrains: Vec<String>,
    pub features: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    #[serde_as(as = "OneOrMany<_>")]
    pub track_features: Vec<String>,
    pub license: Option<String>,
    pub license_family: Option<String>,
    pub noarch: Option<NoArchType>,
    pub python_site_packages_path: Option<String>,
    pub size: Option<u64>,
    #[serde_as(as = "Option<crate::utils::serde::Timestamp>")]
    pub timestamp: Option<chrono::DateTime<chrono::Utc>>,
    #[serde(default, skip_serializing_if = "BTreeSet::is_empty")]
    pub purls: BTreeSet<PackageUrl>,
}

fn create_platforms_and_packages(
    mut per_package: ahash::HashMap<rattler_conda_types::Platform, IndexSet<PackageIndex>>,
    packages: &[LockedPackage],
) -> Result<
    (
        Vec<PlatformData>,
        ahash::HashMap<PlatformIndex, EnvironmentPackages>,
    ),
    ParseCondaLockError,
> {
    let mut unique_platforms = ahash::HashSet::default();

    let mut platforms = per_package
        .keys()
        .filter(|platform| unique_platforms.insert(**platform))
        .map(|platform| PlatformData {
            name: PlatformName::try_from(platform.to_string())
                .expect("All `rattler_conda_types::Platform` are valid platform names"),
            subdir: *platform,
            virtual_packages: Vec::new(),
        })
        .collect::<Vec<_>>();

    platforms.sort_by_key(|p| p.name.to_string());

    let environment_packages = per_package
        .drain()
        .map(|(platform, indices)| {
            Ok((
                PlatformIndex(
                    platforms
                        .iter()
                        .position(|p| p.name.as_str() == platform.as_str())
                        .expect("All Platforms in this hashmap were added before"),
                ),
                EnvironmentPackages::from_indices(indices, packages)?,
            ))
        })
        .collect::<Result<_, ParseCondaLockError>>()?;
    Ok((platforms, environment_packages))
}

/// A function that enables parsing of lock files version 3 or lower.
pub fn parse_v3_or_lower(
    document: serde_yaml::Value,
    version: FileFormatVersion,
) -> Result<LockFile, ParseCondaLockError> {
    let lock_file: LockFileV3 =
        serde_yaml::from_value(document).map_err(ParseCondaLockError::ParseError)?;

    // Iterate over all packages, deduplicate them and store the list of packages
    // per platform. There might be duplicates for noarch packages.
    // let mut conda_packages = IndexSet::with_capacity(lock_file.package.len());
    // let mut pypi_packages = IndexSet::with_capacity(lock_file.package.len());
    let mut packages = IndexSet::with_capacity(lock_file.package.len());

    let mut per_platform: ahash::HashMap<rattler_conda_types::Platform, IndexSet<PackageIndex>> =
        ahash::HashMap::default();
    for package in lock_file.package {
        let LockedPackageV3 { platform, kind } = package;

        let pkg: PackageIndex = match kind {
            LockedPackageKindV3::Conda(value) => {
                let md5 = match value.hash {
                    PackageHashes::Md5(md5) | PackageHashes::Md5Sha256(md5, _) => Some(md5),
                    PackageHashes::Sha256(_) => None,
                };
                let sha256 = match value.hash {
                    PackageHashes::Sha256(sha256) | PackageHashes::Md5Sha256(_, sha256) => {
                        Some(sha256)
                    }
                    PackageHashes::Md5(_) => None,
                };

                // Guess the subdir from the url.
                let subdir = value
                    .subdir
                    .as_deref()
                    .or_else(|| {
                        value
                            .url
                            .path_segments()
                            .and_then(|split| split.rev().nth(1))
                    })
                    .and_then(|subdir_str| rattler_conda_types::Platform::from_str(subdir_str).ok())
                    .unwrap_or(platform);

                let location = UrlOrPath::Url(value.url).normalize().into_owned();
                let derived = LocationDerivedFields::new(&location);

                let build = value
                    .build
                    .or_else(|| derived.build.clone())
                    .unwrap_or_default();
                let build_number = value
                    .build_number
                    .or_else(|| derive_build_number_from_build(&build))
                    .unwrap_or(0);
                let derived_noarch = derive_noarch_type(
                    derived.subdir.as_deref().unwrap_or(subdir.as_str()),
                    derived.build.as_deref().unwrap_or(&build),
                );
                let (derived_arch, derived_platform) =
                    derive_arch_and_platform(derived.subdir.as_deref().unwrap_or(subdir.as_str()));

                PackageIndex(
                    packages
                        .insert_full(LockedPackage::Conda(
                            LegacyCondaPackageData::Binary(LegacyCondaBinaryData {
                                channel: derived.channel,
                                file_name: derived.identifier.map_or_else(
                                    || format!("{}-{}-{}.conda", value.name, value.version, build),
                                    |id| id.to_file_name(),
                                ),
                                package_record: PackageRecord {
                                    arch: value.arch.or(derived_arch),
                                    build,
                                    build_number,
                                    constrains: value.constrains,
                                    depends: value.dependencies,
                                    experimental_extra_depends: std::collections::BTreeMap::new(),
                                    features: value.features,
                                    flags: Vec::new(),
                                    legacy_bz2_md5: None,
                                    legacy_bz2_size: None,
                                    license: value.license,
                                    license_family: value.license_family,
                                    md5,
                                    name: PackageName::new_unchecked(value.name),
                                    noarch: value.noarch.unwrap_or(derived_noarch),
                                    platform: derived_platform,
                                    sha256,
                                    size: value.size,
                                    subdir: subdir.to_string(),
                                    timestamp: value.timestamp.map(Into::into),
                                    track_features: value.track_features,
                                    version: value.version,
                                    purls: value.purls.is_empty().not().then_some(value.purls),
                                    python_site_packages_path: value.python_site_packages_path,
                                    run_exports: None,
                                },
                                location,
                            })
                            .into(),
                        ))
                        .0,
                )
            }
            LockedPackageKindV3::Pypi(pkg) => PackageIndex(
                packages
                    .insert_full(LockedPackage::Pypi(PypiPackageData::from(
                        PypiDistributionData {
                            name: pep508_rs::PackageName::new(pkg.name)?,
                            version: pkg.version,
                            requires_dist: pkg.requires_dist,
                            requires_python: pkg.requires_python,
                            location: Verbatim::new(UrlOrPath::Url(pkg.url)),
                            hash: pkg.hash,
                            index_url: None,
                        },
                    )))
                    .0,
            ),
        };

        per_platform
            .entry(package.platform)
            .or_default()
            .insert(pkg);
    }

    let packages: Vec<LockedPackage> = packages.into_iter().collect();
    let (platforms, environment_packages) = create_platforms_and_packages(per_platform, &packages)?;
    // Construct the default environment
    let default_environment = EnvironmentData {
        channels: lock_file.metadata.channels,
        indexes: None,
        packages: environment_packages,
        options: SolveOptions::default(),
    };

    Ok(LockFile {
        inner: Arc::new(LockFileInner {
            version,
            platforms,
            packages,
            environment_lookup: [(DEFAULT_ENVIRONMENT_NAME.to_string(), EnvironmentIndex(0))]
                .into_iter()
                .collect(),
            environments: vec![default_environment],
        }),
    })
}