Support alternative registries.

cargo-add gains a new registry flag.
cargo-upgrade gets any alternative registry URLs from cargo metadata.

Cargo.toml

@@ -68,6 +68,7 @@ git2 = "0.9.2"
 hex = "0.3.2"
 url = "2.1.0"
 toml = "0.5.1"
+itertools = "0.8.0"
 
 [dependencies.semver]
 features = ["serde"]

src/bin/add/args.rs

@@ -1,6 +1,6 @@
 //! Handle `cargo add` arguments
 
-use cargo_edit::{find, Dependency};
+use cargo_edit::{find, registry_url, Dependency};
 use cargo_edit::{get_latest_dependency, CrateName};
 use semver;
 use std::path::PathBuf;
@@ -104,6 +104,10 @@ pub struct Args {
     /// Keep dependencies sorted
     #[structopt(long = "sort", short = "s")]
     pub sort: bool,
+
+    /// Registry to use
+    #[structopt(long = "registry", conflicts_with = "git", conflicts_with = "path")]
+    pub registry: Option<String>,
 }
 
 fn parse_version_req(s: &str) -> Result<&str> {
@@ -153,6 +157,8 @@ impl Args {
         } else {
             assert_eq!(self.git.is_some() && self.vers.is_some(), false);
             assert_eq!(self.git.is_some() && self.path.is_some(), false);
+            assert_eq!(self.git.is_some() && self.registry.is_some(), false);
+            assert_eq!(self.path.is_some() && self.registry.is_some(), false);
 
             let mut dependency = Dependency::new(crate_name.name());
 
@@ -165,12 +171,18 @@ impl Args {
             if let Some(version) = &self.vers {
                 dependency = dependency.set_version(parse_version_req(version)?);
             }
+            let registry_url = if let Some(registry) = &self.registry {
+                Some(registry_url(&find(&self.manifest_path)?, Some(registry))?)
+            } else {
+                None
+            };
 
             if self.git.is_none() && self.path.is_none() && self.vers.is_none() {
                 let dep = get_latest_dependency(
                     crate_name.name(),
                     self.allow_prerelease,
                     &find(&self.manifest_path)?,
+                    &registry_url,
                 )?;
                 let v = format!(
                     "{prefix}{version}",
@@ -182,6 +194,12 @@ impl Args {
                 dependency = dep.set_version(&v);
             }
 
+            // Set the registry after getting the latest version as
+            // get_latest_dependency returns a registry-less Dependency
+            if let Some(registry) = &self.registry {
+                dependency = dependency.set_registry(registry);
+            }
+
             Ok(dependency)
         }
     }
@@ -236,6 +254,7 @@ impl Default for Args {
             quiet: false,
             offline: true,
             sort: false,
+            registry: None,
         }
     }
 }

src/bin/add/main.rs

@@ -15,12 +15,12 @@
 extern crate error_chain;
 
 use crate::args::{Args, Command};
-use cargo_edit::{find, update_registry_index, Dependency, Manifest};
+use cargo_edit::{find, registry_url, update_registry_index, Dependency, Manifest};
 use std::io::Write;
 use std::process;
 use structopt::StructOpt;
-use toml_edit::Item as TomlItem;
 use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
+use toml_edit::Item as TomlItem;
 
 mod args;
 
@@ -83,7 +83,11 @@ fn handle_add(args: &Args) -> Result<()> {
     let deps = &args.parse_dependencies()?;
 
     if !args.offline {
-        update_registry_index(&find(manifest_path)?)?;
+        let url = registry_url(
+            &find(&manifest_path)?,
+            args.registry.as_ref().map(String::as_ref),
+        )?;
+        update_registry_index(&url)?;
     }
 
     deps.iter()

src/bin/upgrade/main.rs

@@ -16,15 +16,18 @@ extern crate error_chain;
 
 use crate::errors::*;
 use cargo_edit::{
-    find, get_latest_dependency, update_registry_index, CrateName, Dependency, LocalManifest,
+    find, get_latest_dependency, registry_url, update_registry_index, CrateName, Dependency,
+    LocalManifest,
 };
 use failure::Fail;
+use itertools::Itertools;
 use std::collections::HashMap;
 use std::io::Write;
 use std::path::{Path, PathBuf};
 use std::process;
 use structopt::StructOpt;
 use termcolor::{BufferWriter, Color, ColorChoice, ColorSpec, WriteColor};
+use url::Url;
 
 mod errors {
     error_chain! {
@@ -153,30 +156,44 @@ impl Manifests {
             }
         }
 
-        Ok(DesiredUpgrades(if only_update.is_empty() {
-            // User hasn't asked for any specific dependencies to be upgraded, so upgrade all the
-            // dependencies.
+        // Map the names of user-specified dependencies to the (optionally) requested version.
+        let selected_dependencies = only_update
+            .into_iter()
+            .map(|name| {
+                if let Some(dependency) = CrateName::new(&name.clone()).parse_as_version()? {
+                    Ok((
+                        dependency.name.clone(),
+                        dependency.version().map(String::from),
+                    ))
+                } else {
+                    Ok((name, None))
+                }
+            })
+            .collect::<Result<HashMap<_, _>>>()?;
+
+        Ok(DesiredUpgrades(
             self.0
                 .iter()
                 .flat_map(|&(_, ref package)| package.dependencies.clone())
                 .filter(is_version_dep)
-                .map(|dependency| (dependency.name, None))
-                .collect()
-        } else {
-            only_update
-                .into_iter()
-                .map(|name| {
-                    if let Some(dependency) = CrateName::new(&name.clone()).parse_as_version()? {
-                        Ok((
-                            dependency.name.clone(),
-                            dependency.version().map(String::from),
-                        ))
+                .filter_map(|dependency| {
+                    if selected_dependencies.is_empty() {
+                        // User hasn't asked for any specific dependencies to be upgraded,
+                        // so upgrade all the dependencies.
+                        Some((dependency.name, (dependency.registry, None)))
                     } else {
-                        Ok((name, None))
+                        // User has asked for specific dependencies. Check if this dependency
+                        // was specified, populating the registry from the lockfile metadata.
+                        match selected_dependencies.get(&dependency.name) {
+                            Some(version) => {
+                                Some((dependency.name, (dependency.registry, version.clone())))
+                            }
+                            None => None,
+                        }
                     }
                 })
-                .collect::<Result<_>>()?
-        }))
+                .collect(),
+        ))
     }
 
     /// Upgrade the manifests on disk following the previously-determined upgrade schema.
@@ -211,8 +228,9 @@ impl Manifests {
     }
 }
 
-/// The set of dependencies to be upgraded, alongside desired versions, if specified by the user.
-struct DesiredUpgrades(HashMap<String, Option<String>>);
+/// The set of dependencies to be upgraded, alongside the registries returned from cargo metadata, and
+/// the desired versions, if specified by the user.
+struct DesiredUpgrades(HashMap<String, (Option<String>, Option<String>)>);
 
 /// The complete specification of the upgrades that will be performed. Map of the dependency names
 /// to the new versions.
@@ -224,11 +242,17 @@ impl DesiredUpgrades {
     fn get_upgraded(self, allow_prerelease: bool, manifest_path: &Path) -> Result<ActualUpgrades> {
         self.0
             .into_iter()
-            .map(|(name, version)| {
+            .map(|(name, (registry, version))| {
                 if let Some(v) = version {
                     Ok((name, v))
                 } else {
-                    get_latest_dependency(&name, allow_prerelease, manifest_path)
+                    let registry_url = match registry {
+                        Some(x) => Some(Url::parse(&x).map_err(|_| {
+                            ErrorKind::CargoEditLib(::cargo_edit::ErrorKind::InvalidCargoConfig)
+                        })?),
+                        None => None,
+                    };
+                    get_latest_dependency(&name, allow_prerelease, manifest_path, &registry_url)
                         .map(|new_dep| {
                             (
                                 name,
@@ -259,7 +283,8 @@ fn process(args: Args) -> Result<()> {
     } = args;
 
     if !args.offline {
-        update_registry_index(&find(&manifest_path)?)?;
+        let url = registry_url(&find(&manifest_path)?, None)?;
+        update_registry_index(&url)?;
     }
 
     let manifests = if all {
@@ -270,6 +295,21 @@ fn process(args: Args) -> Result<()> {
 
     let existing_dependencies = manifests.get_dependencies(dependency)?;
 
+    // Update indices for any alternative registries, unless
+    // we're offline.
+    if !args.offline {
+        for registry_url in existing_dependencies
+            .0
+            .values()
+            .filter_map(|(registry, _)| registry.as_ref())
+            .unique()
+        {
+            update_registry_index(&Url::parse(registry_url).map_err(|_| {
+                ErrorKind::CargoEditLib(::cargo_edit::ErrorKind::InvalidCargoConfig)
+            })?)?;
+        }
+    }
+
     let upgraded_dependencies =
         existing_dependencies.get_upgraded(allow_prerelease, &find(&manifest_path)?)?;
 

src/dependency.rs

@@ -5,6 +5,7 @@ enum DependencySource {
     Version {
         version: Option<String>,
         path: Option<String>,
+        registry: Option<String>,
     },
     Git(String),
 }
@@ -28,6 +29,7 @@ impl Default for Dependency {
             source: DependencySource::Version {
                 version: None,
                 path: None,
+                registry: None,
             },
         }
     }
@@ -49,13 +51,14 @@ impl Dependency {
         // ("version requirement […] includes semver metadata which will be ignored")
         let version = version.split('+').next().unwrap();
         let old_source = self.source;
-        let old_path = match old_source {
-            DependencySource::Version { path, .. } => path,
-            _ => None,
+        let (old_path, old_registry) = match old_source {
+            DependencySource::Version { path, registry, .. } => (path, registry),
+            _ => (None, None),
         };
         self.source = DependencySource::Version {
             version: Some(version.into()),
             path: old_path,
+            registry: old_registry,
         };
         self
     }
@@ -76,6 +79,7 @@ impl Dependency {
         self.source = DependencySource::Version {
             version: old_version,
             path: Some(path.into()),
+            registry: None,
         };
         self
     }
@@ -92,6 +96,21 @@ impl Dependency {
         self
     }
 
+    /// Set the value of registry for the dependency
+    pub fn set_registry(mut self, registry: &str) -> Dependency {
+        let old_source = self.source;
+        let old_version = match old_source {
+            DependencySource::Version { version, .. } => version,
+            _ => None,
+        };
+        self.source = DependencySource::Version {
+            version: old_version,
+            path: None,
+            registry: Some(registry.into()),
+        };
+        self
+    }
+
     /// Get version of dependency
     pub fn version(&self) -> Option<&str> {
         if let DependencySource::Version {
@@ -121,20 +140,28 @@ impl Dependency {
                     DependencySource::Version {
                         version: Some(v),
                         path: None,
+                        registry: None,
                     },
                 ) => toml_edit::value(v),
                 // Other cases are represented as an inline table
                 (optional, default_features, source) => {
                     let mut data = toml_edit::InlineTable::default();
 
                     match source {
-                        DependencySource::Version { version, path } => {
+                        DependencySource::Version {
+                            version,
+                            path,
+                            registry,
+                        } => {
                             if let Some(v) = version {
                                 data.get_or_insert("version", v);
                             }
                             if let Some(p) = path {
                                 data.get_or_insert("path", p);
                             }
+                            if let Some(r) = registry {
+                                data.get_or_insert("registry", r);
+                            }
                         }
                         DependencySource::Git(v) => {
                             data.get_or_insert("git", v);

src/errors.rs

@@ -65,5 +65,9 @@ error_chain! {
             description("Unable to find the source specified by 'replace-with'")
             display("The source '{}' could not be found", name)
         }
+        /// Unable to find the specified registry
+        NoSuchRegistryFound(name: String) {
+            display("The registry '{}' could not be found", name)
+        }
     }
 }