(wip) Add --prerelease dispatch to uv

astral-sh · Mar 3, 2024 · 835c37b · 835c37b
1 parent 2111778
commit 835c37b
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 8 deletions.
diff --git a/rye/src/cli/add.rs b/rye/src/cli/add.rs
@@ -206,7 +206,13 @@ pub struct Args {
     #[arg(long, conflicts_with = "dev", conflicts_with = "excluded")]
     optional: Option<String>,
     /// Include pre-releases when finding a package version.
-    #[arg(long)]
+    #[arg(
+        long,
+        default_value = "if-necessary-or-explicit",
+        conflicts_with = "pre"
+    )]
+    prerelease: Option<String>,
+    #[clap(long, hide = true, conflicts_with = "prerelease")]
     pre: bool,
     /// Overrides the pin operator
     #[arg(long)]
@@ -260,13 +266,18 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
 
     if !cmd.excluded {
         if cfg.use_uv() {
+            let pre = if cmd.pre {
+                Some("allow".to_string())
+            } else {
+                cmd.prerelease
+            };
             sync(SyncOptions::python_only().pyproject(None))
                 .context("failed to sync ahead of add")?;
             resolve_requirements_with_uv(
                 &pyproject_toml,
                 &py_ver,
                 &mut requirements,
-                cmd.pre,
+                pre.as_ref(),
                 output,
                 &default_operator,
             )?;
@@ -444,7 +455,7 @@ fn resolve_requirements_with_uv(
     pyproject_toml: &PyProject,
     py_ver: &PythonVersion,
     requirements: &mut [Requirement],
-    pre: bool,
+    prerelease: Option<&String>,
     output: CommandOutput,
     default_operator: &Operator,
 ) -> Result<(), Error> {
@@ -458,8 +469,8 @@ fn resolve_requirements_with_uv(
             .arg("--no-header")
             .arg("-")
             .env("VIRTUAL_ENV", pyproject_toml.venv_path().as_os_str());
-        if pre {
-            cmd.arg("--prerelease=allow");
+        if let Some(it) = prerelease {
+            cmd.arg("--prerelease").arg(it);
         }
         if output == CommandOutput::Quiet {
             cmd.arg("-q");

diff --git a/rye/src/cli/lock.rs b/rye/src/cli/lock.rs
@@ -23,7 +23,13 @@ pub struct Args {
     #[arg(long)]
     update_all: bool,
     /// Update to pre-release versions
-    #[arg(long)]
+    #[arg(
+        long,
+        default_value = "if-necessary-or-explicit",
+        conflicts_with = "pre"
+    )]
+    prerelease: Option<String>,
+    #[clap(long, hide = true, conflicts_with = "prerelease")]
     pre: bool,
     /// Extras/features to enable when locking the workspace.
     #[arg(long)]
@@ -50,6 +56,7 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
         lock_options: LockOptions {
             update: cmd.update,
             update_all: cmd.update_all,
+            prerelease: cmd.prerelease,
             pre: cmd.pre,
             features: cmd.features,
             all_features: cmd.all_features,

diff --git a/rye/src/cli/sync.rs b/rye/src/cli/sync.rs
@@ -31,8 +31,14 @@ pub struct Args {
     /// Update all packages to the latest
     #[arg(long)]
     update_all: bool,
-    /// Update to pre-release versions
-    #[arg(long)]
+    /// Update to pre-release versions (uv)
+    #[arg(
+        long,
+        default_value = "if-necessary-or-explicit",
+        conflicts_with = "pre"
+    )]
+    prerelease: Option<String>,
+    #[clap(long, hide = true, conflicts_with = "prerelease")]
     pre: bool,
     /// Extras/features to enable when synching the workspace.
     #[arg(long)]
@@ -66,6 +72,7 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
         lock_options: LockOptions {
             update: cmd.update,
             update_all: cmd.update_all,
+            prerelease: cmd.prerelease,
             pre: cmd.pre,
             features: cmd.features,
             all_features: cmd.all_features,

diff --git a/rye/src/lock.rs b/rye/src/lock.rs
@@ -66,6 +66,8 @@ pub struct LockOptions {
     pub update_all: bool,
     /// Update specific packages.
     pub update: Vec<String>,
+    /// Pick pre-release versions (uv).
+    pub prerelease: Option<String>,
     /// Pick pre-release versions.
     pub pre: bool,
     /// A list of features (extras) to enable when locking
@@ -102,6 +104,7 @@ impl LockOptions {
             if let Some(m) = PARAM_RE.captures(line) {
                 let value = &m[2];
                 match &m[1] {
+                    // TODO(cnpryer)
                     "pre" => rv.pre = rv.pre || serde_json::from_str(value)?,
                     "features" => {
                         if rv.features.is_empty() {
@@ -391,6 +394,8 @@ fn generate_lockfile(
 
         if lock_options.pre {
             cmd.arg("--prerelease=allow");
+        } else if let Some(it) = lock_options.prerelease.as_ref() {
+            cmd.arg("--prerelease").arg(it);
         }
         // this primarily exists for testing
         if let Ok(dt) = env::var("__RYE_UV_EXCLUDE_NEWER") {