Support uv --override option (#668)

astral-sh · Apr 20, 2024 · a019cd7 · a019cd7
1 parent 5440506
commit a019cd7
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 10 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,8 @@ that were not yet released.
 
 _Unreleased_
 
+- Support dependency overrides via `tool.rye.override-dependencies` when using uv. #668
+
 <!-- released start -->
 
 ## 0.32.0

diff --git a/rye/src/cli/test.rs b/rye/src/cli/test.rs
@@ -137,8 +137,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" {

diff --git a/rye/src/lock.rs b/rye/src/lock.rs
@@ -178,12 +178,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,
@@ -194,6 +197,20 @@ 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,
@@ -263,10 +280,10 @@ 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())
@@ -285,7 +302,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
@@ -334,23 +351,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,
@@ -367,6 +387,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,
@@ -405,12 +426,16 @@ fn generate_lockfile(
             .lockfile(
                 py_ver,
                 requirements_file_in,
+                overrides_file_in,
                 &requirements_file,
                 lock_options.pre,
                 env::var("__RYE_UV_EXCLUDE_NEWER").ok(),
                 upgrade,
             )?;
     } else {
+        if overrides_file_in.is_some() {
+            bail!("dependency overrides are only supported by uv");
+        }
         let mut cmd = Command::new(get_pip_compile(py_ver, output)?);
         // legacy pip tools requires some extra parameters
         if get_pip_tools_version(py_ver) == PipToolsVersion::Legacy {

diff --git a/rye/src/pyproject.rs b/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("overrides"),
             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")

diff --git a/rye/src/uv.rs b/rye/src/uv.rs
@@ -307,6 +307,7 @@ impl Uv {
         &self,
         py_version: &PythonVersion,
         source: &Path,
+        overrides: Option<&Path>,
         target: &Path,
         allow_prerelease: bool,
         exclude_newer: Option<String>,
@@ -333,6 +334,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 {}",