From 68280854d2590a5ea4f59cb5f03de4dfede8008a Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Wed, 21 Jun 2023 00:50:04 +0200
Subject: [PATCH] Added --pin argument to rye add

---
 CHANGELOG.md       |  2 ++
 rye/src/cli/add.rs | 37 +++++++++++++++++++++++++++++++------
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8858046f6c..bd562823de 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@ that were not yet released.
 
 _Unreleased_
 
+- `rye add` now accepts `--pin` to let one override the type of pin to use.  #341
+
 - Added `rye config` to read and manipulate the `config.toml` file.  #339
 
 - Added support for the new `behavior.global-python` flag which turns on global
diff --git a/rye/src/cli/add.rs b/rye/src/cli/add.rs
index 0a782aecd6..e98407097c 100644
--- a/rye/src/cli/add.rs
+++ b/rye/src/cli/add.rs
@@ -4,7 +4,7 @@ use std::process::{Command, Stdio};
 use std::str::FromStr;
 
 use anyhow::{anyhow, bail, Context, Error};
-use clap::Parser;
+use clap::{Parser, ValueEnum};
 use console::style;
 use pep440_rs::{Operator, Version, VersionSpecifier, VersionSpecifiers};
 use pep508_rs::{Requirement, VersionOrUrl};
@@ -88,6 +88,26 @@ pub struct ReqExtras {
     features: Vec<String>,
 }
 
+#[derive(ValueEnum, Copy, Clone, Debug, PartialEq)]
+enum Pin {
+    #[value(alias = "exact", alias = "==", alias = "eq")]
+    Equal,
+    #[value(alias = "tilde", alias = "compatible", alias = "~=")]
+    TildeEqual,
+    #[value(alias = ">=", alias = "ge", alias = "gte")]
+    GreaterThanEqual,
+}
+
+impl From<Pin> for Operator {
+    fn from(value: Pin) -> Self {
+        match value {
+            Pin::Equal => Operator::Equal,
+            Pin::TildeEqual => Operator::TildeEqual,
+            Pin::GreaterThanEqual => Operator::GreaterThanEqual,
+        }
+    }
+}
+
 impl ReqExtras {
     pub fn force_absolute(&mut self) {
         self.absolute = true;
@@ -178,6 +198,9 @@ pub struct Args {
     /// Include pre-releases when finding a package version.
     #[arg(long)]
     pre: bool,
+    /// Overrides the pin operator
+    #[arg(long)]
+    pin: Option<Pin>,
     /// Enables verbose diagnostics.
     #[arg(short, long)]
     verbose: bool,
@@ -207,7 +230,10 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
     } else {
         DependencyKind::Normal
     };
-    let default_operator = Config::current().default_dependency_operator();
+    let default_operator = match cmd.pin {
+        Some(pin) => Operator::from(pin),
+        None => Config::current().default_dependency_operator(),
+    };
 
     for str_requirement in cmd.requirements {
         let mut requirement = Requirement::from_str(&str_requirement)?;
@@ -290,10 +316,9 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
                             // local versions or versions with only one component cannot
                             // use ~= but need to use ==.
                             match default_operator {
-                                Operator::EqualStar
-                                    if version.is_local() || version.release.len() < 2 =>
-                                {
-                                    Operator::TildeEqual
+                                _ if version.is_local() => Operator::Equal,
+                                Operator::TildeEqual if version.release.len() < 2 => {
+                                    Operator::GreaterThanEqual
                                 }
                                 ref other => other.clone(),
                             },