Support uv --override option (astral-sh#668)

rye/src/cli/add.rs

@@ -203,11 +203,32 @@ pub struct Args {
     #[arg(short, long)]
     dev: bool,
     /// Add this as an excluded dependency that will not be installed even if it's a sub dependency.
-    #[arg(long, conflicts_with = "dev", conflicts_with = "optional")]
+    #[arg(
+        long,
+        conflicts_with = "dev",
+        conflicts_with = "optional",
+        conflicts_with = "override"
+    )]
     excluded: bool,
     /// Add this to an optional dependency group.
-    #[arg(long, conflicts_with = "dev", conflicts_with = "excluded")]
+    #[arg(
+        long,
+        conflicts_with = "dev",
+        conflicts_with = "excluded",
+        conflicts_with = "override"
+    )]
     optional: Option<String>,
+    /// Add this as an override dependency.
+    #[arg(
+        long,
+        conflicts_with = "dev",
+        conflicts_with = "optional",
+        conflicts_with = "excluded"
+    )]
+    r#override: bool,
+    /// Include pre-releases when finding a package version.
+    #[arg(long)]
+    pre: bool,
     /// Overrides the pin operator
     #[arg(long)]
     pin: Option<Pin>,
@@ -252,6 +273,8 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
         DependencyKind::Excluded
     } else if let Some(ref section) = cmd.optional {
         DependencyKind::Optional(section.into())
+    } else if cmd.r#override {
+        DependencyKind::Override
     } else {
         DependencyKind::Normal
     };

rye/src/cli/test.rs

@@ -158,8 +158,8 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
 fn has_pytest_dependency(projects: &[PyProject]) -> Result<bool, Error> {
     for project in projects {
         for dep in project
-            .iter_dependencies(DependencyKind::Dev)
-            .chain(project.iter_dependencies(DependencyKind::Normal))
+            .iter_dependencies(&DependencyKind::Dev)
+            .chain(project.iter_dependencies(&DependencyKind::Normal))
         {
             if let Ok(req) = dep.expand(|name| std::env::var(name).ok()) {
                 if normalize_package_name(&req.name) == "pytest" {

rye/src/lock.rs

@@ -196,12 +196,15 @@ pub fn update_workspace_lockfile(
 
     req_file.flush()?;
 
-    let exclusions = find_exclusions(&projects)?;
+    let exclusions = find_requirements(&projects, &DependencyKind::Excluded)?;
+    let overrides = find_requirements(&projects, &DependencyKind::Override)?;
+    let overrides_file = maybe_write_requirements_to_temp(&overrides)?;
     generate_lockfile(
         output,
         py_ver,
         &workspace.path(),
         req_file.path(),
+        overrides_file.as_ref().map(|v| v.path()),
         lockfile,
         sources,
         &lock_options,
@@ -213,6 +216,21 @@ pub fn update_workspace_lockfile(
     Ok(())
 }
 
+fn maybe_write_requirements_to_temp(
+    requirements: &HashSet<Requirement>,
+) -> Result<Option<NamedTempFile>, Error> {
+    if requirements.is_empty() {
+        Ok(None)
+    } else {
+        let mut nt_file = NamedTempFile::new()?;
+        for dep in requirements {
+            writeln!(&nt_file, "{}", dep)?;
+        }
+        nt_file.flush()?;
+        Ok(Some(nt_file))
+    }
+}
+
 /// Tries to restore the lock options from the given lockfile.
 fn restore_lock_options<'o>(
     lockfile: &Path,
@@ -282,10 +300,13 @@ fn collect_workspace_features(
     Some(features_by_project)
 }
 
-fn find_exclusions(projects: &[PyProject]) -> Result<HashSet<Requirement>, Error> {
+fn find_requirements(
+    projects: &[PyProject],
+    kind: &DependencyKind,
+) -> Result<HashSet<Requirement>, Error> {
     let mut rv = HashSet::new();
     for project in projects {
-        for dep in project.iter_dependencies(DependencyKind::Excluded) {
+        for dep in project.iter_dependencies(kind) {
             rv.insert(dep.expand(|name: &str| {
                 if name == "PROJECT_ROOT" {
                     Some(project.workspace_path().to_string_lossy().to_string())
@@ -304,7 +325,7 @@ fn dump_dependencies(
     out: &mut fs::File,
     dep_kind: DependencyKind,
 ) -> Result<(), Error> {
-    for dep in pyproject.iter_dependencies(dep_kind) {
+    for dep in pyproject.iter_dependencies(&dep_kind) {
         if let Ok(expanded_dep) = dep.expand(|_| {
             // we actually do not care what it expands to much, for as long
             // as the end result parses
@@ -355,23 +376,26 @@ pub fn update_single_project_lockfile(
         )?;
     }
 
-    for dep in pyproject.iter_dependencies(DependencyKind::Normal) {
+    for dep in pyproject.iter_dependencies(&DependencyKind::Normal) {
         writeln!(req_file, "{}", dep)?;
     }
     if lock_mode == LockMode::Dev {
-        for dep in pyproject.iter_dependencies(DependencyKind::Dev) {
+        for dep in pyproject.iter_dependencies(&DependencyKind::Dev) {
             writeln!(req_file, "{}", dep)?;
         }
     }
 
     req_file.flush()?;
 
-    let exclusions = find_exclusions(std::slice::from_ref(pyproject))?;
+    let exclusions = find_requirements(std::slice::from_ref(pyproject), &DependencyKind::Excluded)?;
+    let overrides = find_requirements(std::slice::from_ref(pyproject), &DependencyKind::Override)?;
+    let overrides_file = maybe_write_requirements_to_temp(&overrides)?;
     generate_lockfile(
         output,
         py_ver,
         &pyproject.workspace_path(),
         req_file.path(),
+        overrides_file.as_ref().map(|v| v.path()),
         lockfile,
         sources,
         &lock_options,
@@ -389,6 +413,7 @@ fn generate_lockfile(
     py_ver: &PythonVersion,
     workspace_path: &Path,
     requirements_file_in: &Path,
+    overrides_file_in: Option<&Path>,
     lockfile: &Path,
     sources: &ExpandedSources,
     lock_options: &LockOptions,
@@ -428,6 +453,7 @@ fn generate_lockfile(
             .lockfile(
                 py_ver,
                 requirements_file_in,
+                overrides_file_in,
                 &requirements_file,
                 lock_options.pre,
                 env::var("__RYE_UV_EXCLUDE_NEWER").ok(),
@@ -436,6 +462,9 @@ fn generate_lockfile(
                 lock_options.generate_hashes,
             )?;
     } else {
+        if overrides_file_in.is_some() {
+            bail!("dependency overrides require the uv backend");
+        }
         if keyring_provider != KeyringProvider::Disabled {
             bail!("`--keyring-provider` option requires the uv backend");
         }

rye/src/pyproject.rs

@@ -55,6 +55,7 @@ pub enum DependencyKind<'a> {
     Normal,
     Dev,
     Excluded,
+    Override,
     Optional(Cow<'a, str>),
 }
 
@@ -64,6 +65,7 @@ impl<'a> fmt::Display for DependencyKind<'a> {
             DependencyKind::Normal => f.write_str("regular"),
             DependencyKind::Dev => f.write_str("dev"),
             DependencyKind::Excluded => f.write_str("excluded"),
+            DependencyKind::Override => f.write_str("override"),
             DependencyKind::Optional(ref sect) => write!(f, "optional ({})", sect),
         }
     }
@@ -903,6 +905,7 @@ impl PyProject {
             DependencyKind::Normal => &mut self.doc["project"]["dependencies"],
             DependencyKind::Dev => &mut self.doc["tool"]["rye"]["dev-dependencies"],
             DependencyKind::Excluded => &mut self.doc["tool"]["rye"]["excluded-dependencies"],
+            DependencyKind::Override => &mut self.doc["tool"]["rye"]["override-dependencies"],
             DependencyKind::Optional(ref section) => {
                 // add this as a proper non-inline table if it's missing
                 let table = &mut self.doc["project"]["optional-dependencies"];
@@ -934,6 +937,7 @@ impl PyProject {
             DependencyKind::Normal => &mut self.doc["project"]["dependencies"],
             DependencyKind::Dev => &mut self.doc["tool"]["rye"]["dev-dependencies"],
             DependencyKind::Excluded => &mut self.doc["tool"]["rye"]["excluded-dependencies"],
+            DependencyKind::Override => &mut self.doc["tool"]["rye"]["override-dependencies"],
             DependencyKind::Optional(ref section) => {
                 &mut self.doc["project"]["optional-dependencies"][section as &str]
             }
@@ -953,7 +957,7 @@ impl PyProject {
     /// Iterates over all dependencies.
     pub fn iter_dependencies(
         &self,
-        kind: DependencyKind,
+        kind: &DependencyKind,
     ) -> impl Iterator<Item = DependencyRef> + '_ {
         let sec = match kind {
             DependencyKind::Normal => self.doc.get("project").and_then(|x| x.get("dependencies")),
@@ -967,6 +971,11 @@ impl PyProject {
                 .get("tool")
                 .and_then(|x| x.get("rye"))
                 .and_then(|x| x.get("excluded-dependencies")),
+            DependencyKind::Override => self
+                .doc
+                .get("tool")
+                .and_then(|x| x.get("rye"))
+                .and_then(|x| x.get("override-dependencies")),
             DependencyKind::Optional(ref section) => self
                 .doc
                 .get("project")

rye/src/uv.rs

@@ -331,6 +331,7 @@ impl Uv {
         &self,
         py_version: &PythonVersion,
         source: &Path,
+        overrides: Option<&Path>,
         target: &Path,
         allow_prerelease: bool,
         exclude_newer: Option<String>,
@@ -361,6 +362,8 @@ impl Uv {
 
         cmd.arg(source);
 
+        overrides.map(|ref value| cmd.arg("--override").arg(value));
+
         let status = cmd.status().with_context(|| {
             format!(
                 "Unable to run uv pip compile and generate {}",

rye/tests/test_sync.rs

@@ -263,3 +263,50 @@ fn test_autosync_remember() {
     werkzeug==3.0.1
     "###);
 }
+
+#[test]
+fn test_overrides() {
+    // enforce werkzeug==2.3.8 when flask==3.0.0 requires Werkzeug>=3.0.0
+
+    let space = Space::new();
+    space.init("my-project");
+
+    rye_cmd_snapshot!(space.rye_cmd().arg("add").arg("werkzeug==2.3.8").arg("--override").arg("--no-sync"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    Initializing new virtualenv in [TEMP_PATH]/project/.venv
+    Python version: [email protected]
+    Added werkzeug==2.3.8 as override dependency
+
+    ----- stderr -----
+    "###);
+
+    rye_cmd_snapshot!(space.rye_cmd().arg("add").arg("flask==3.0.0").arg("colorama==0.4.6"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    Added flask==3.0.0 as regular dependency
+    Added colorama==0.4.6 as regular dependency
+    Reusing already existing virtualenv
+    Generating production lockfile: [TEMP_PATH]/project/requirements.lock
+    Generating dev lockfile: [TEMP_PATH]/project/requirements-dev.lock
+    Installing dependencies
+    Done!
+
+    ----- stderr -----
+    Built 1 editable in [EXECUTION_TIME]
+    Resolved 8 packages in [EXECUTION_TIME]
+    Downloaded 8 packages in [EXECUTION_TIME]
+    Installed 9 packages in [EXECUTION_TIME]
+     + blinker==1.7.0
+     + click==8.1.7
+     + colorama==0.4.6
+     + flask==3.0.0
+     + itsdangerous==2.1.2
+     + jinja2==3.1.2
+     + markupsafe==2.1.3
+     + my-project==0.1.0 (from file:[TEMP_PATH]/project)
+     + werkzeug==2.3.8
+    "###);
+}