download pypi resolved urls from mirrored urls to improve downloading…

… proformance

src/environment/pypi_prefix.rs

@@ -13,6 +13,7 @@ use std::collections::HashMap;
 use uv_distribution_types::{InstalledDist, Name};
 
 use crate::{install_pypi, lock_file::UvResolutionContext, prefix::Prefix};
+use url::Url;
 
 use super::PythonStatus;
 
@@ -83,6 +84,7 @@ pub async fn update_prefix_pypi(
     platform: Platform,
     non_isolated_packages: Option<Vec<String>>,
     no_build: &pixi_manifest::pypi::pypi_options::NoBuild,
+    mirror_map: &std::collections::HashMap<Url, Vec<Url>>,
 ) -> miette::Result<()> {
     // If we have changed interpreter, we need to uninstall all site-packages from
     // the old interpreter We need to do this before the pypi prefix update,
@@ -149,6 +151,7 @@ pub async fn update_prefix_pypi(
                 platform,
                 non_isolated_packages,
                 no_build,
+                mirror_map,
             )
         },
     )

src/install_pypi/conversions.rs

@@ -84,12 +84,15 @@ pub enum ConvertToUvDistError {
 
     #[error(transparent)]
     UvPepTypes(#[from] ConversionError),
+    #[error(transparent)]
+    ParseError(#[from] url::ParseError),
 }
 
 /// Convert from a PypiPackageData to a uv [`distribution_types::Dist`]
 pub fn convert_to_dist(
     pkg: &PypiPackageData,
     lock_file_dir: &Path,
+    mirror_map: &std::collections::HashMap<Url, Vec<Url>>,
 ) -> Result<Dist, ConvertToUvDistError> {
     // Figure out if it is a url from the registry or a direct url
     let dist = match &pkg.location {
@@ -125,6 +128,21 @@ pub fn convert_to_dist(
             }
         }
         UrlOrPath::Url(url) => {
+            // try to find mapped url
+            let mut mapped_url: Option<Url> = None;
+            'outer: for (url_prefix, mapped_prefixs) in mirror_map {
+                if url.as_str().starts_with(url_prefix.as_str()) {
+                    for mapped_prefix in mapped_prefixs {
+                        if ! is_direct_url(mapped_prefix.scheme()) {
+                            let new_url_str = url.as_str().replacen(url_prefix.as_str(), mapped_prefix.as_str(), 1);
+                            let new_url = Url::parse(&new_url_str)?;
+                            mapped_url = Some(new_url);
+                            break 'outer;
+                        }
+                    }
+                }
+            }
+
             // We consider it to be a registry url
             // Extract last component from registry url
             // should be something like `package-0.1.0-py3-none-any.whl`
@@ -141,7 +159,7 @@ pub fn convert_to_dist(
             // Now we can convert the locked data to a [`distribution_types::File`]
             // which is essentially the file information for a wheel or sdist
             let file = locked_data_to_file(
-                url,
+                mapped_url.as_ref().unwrap_or(url),
                 pkg.hash.as_ref(),
                 filename_decoded.as_ref(),
                 pkg.requires_python.clone(),
@@ -211,6 +229,7 @@ pub fn convert_to_dist(
         }
     };
 
+    tracing::info!("dist={}", dist);
     Ok(dist)
 }
 
@@ -243,7 +262,7 @@ mod tests {
 
         // Convert the locked data to a uv dist
         // check if it does not panic
-        let dist = convert_to_dist(&locked, &PathBuf::new())
+        let dist = convert_to_dist(&locked, &PathBuf::new(), &std::collections::HashMap<Url, Vec<Url>>::new())
             .expect("could not convert wheel with special chars to dist");
 
         // Check if the dist is a built dist

src/install_pypi/mod.rs

@@ -26,6 +26,8 @@ use uv_python::{Interpreter, PythonEnvironment};
 use uv_resolver::FlatIndex;
 use uv_types::HashStrategy;
 
+use url::Url;
+
 use crate::{
     lock_file::UvResolutionContext,
     prefix::Prefix,
@@ -57,6 +59,7 @@ pub async fn update_python_distributions(
     platform: Platform,
     non_isolated_packages: Option<Vec<String>>,
     no_build: &pixi_manifest::pypi::pypi_options::NoBuild,
+    mirror_map: &std::collections::HashMap<Url, Vec<Url>>,
 ) -> miette::Result<()> {
     let start = std::time::Instant::now();
 
@@ -190,7 +193,7 @@ pub async fn update_python_distributions(
         reinstalls,
         extraneous,
     } = InstallPlanner::new(uv_context.cache.clone(), lock_file_dir)
-        .plan(&site_packages, registry_index, &required_map)
+        .plan(&site_packages, registry_index, &required_map, mirror_map)
         .into_diagnostic()
         .context("error while determining PyPI installation plan")?;
 

src/install_pypi/plan/planner.rs

@@ -21,6 +21,8 @@ use super::{
     InstallReason, NeedReinstall, PyPIInstallPlan,
 };
 
+use url::Url;
+
 /// Struct that handles the planning of the installation
 /// of the PyPI packages into an existing conda environment with specific
 /// locked data
@@ -75,6 +77,7 @@ impl InstallPlanner {
         remote: &mut Vec<(Dist, InstallReason)>,
         dist_cache: &mut impl CachedDistProvider<'a>,
         op_to_reason: Op,
+        mirror_map: &std::collections::HashMap<Url, Vec<Url>>,
     ) -> Result<(), InstallPlannerError> {
         // Okay so we need to re-install the package
         // let's see if we need the remote or local version
@@ -83,7 +86,7 @@ impl InstallPlanner {
         // then we should get it from the remote
         if self.uv_cache.must_revalidate(name) {
             remote.push((
-                convert_to_dist(required_pkg, &self.lock_file_dir)?,
+                convert_to_dist(required_pkg, &self.lock_file_dir, mirror_map)?,
                 op_to_reason.stale(),
             ));
             return Ok(());
@@ -97,7 +100,7 @@ impl InstallPlanner {
         // If we don't have it in the cache we need to download it
         } else {
             remote.push((
-                convert_to_dist(required_pkg, &self.lock_file_dir)?,
+                convert_to_dist(required_pkg, &self.lock_file_dir, mirror_map)?,
                 op_to_reason.missing(),
             ));
         }
@@ -116,6 +119,7 @@ impl InstallPlanner {
         site_packages: &'a Installed,
         mut dist_cache: Cached,
         required_pkgs: &'a HashMap<uv_normalize::PackageName, &PypiPackageData>,
+        mirror_map: &std::collections::HashMap<Url, Vec<Url>>,
     ) -> Result<PyPIInstallPlan, InstallPlannerError> {
         // Packages to be removed
         let mut extraneous = vec![];
@@ -184,6 +188,7 @@ impl InstallPlanner {
                         &mut remote,
                         &mut dist_cache,
                         reasons::Reinstall,
+                        mirror_map,
                     )?;
                 }
                 // Second case we are not managing the package
@@ -214,6 +219,7 @@ impl InstallPlanner {
                 &mut remote,
                 &mut dist_cache,
                 reasons::Install,
+                mirror_map,
             )?;
         }
 

src/lock_file/update.rs

@@ -440,6 +440,7 @@ impl<'p> LockFileDerivedData<'p> {
             environment.best_platform(),
             non_isolated_packages,
             &no_build,
+            &self.workspace.config().mirror_map(),
         )
         .await
         .with_context(|| {