From 582c94cec313d91923dbd81666c0462fe2de72a4 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Mon, 13 Nov 2023 12:17:01 -0800
Subject: [PATCH] Add missing-dot fix to lenient requirements (#416)

Part of https://github.com/astral-sh/puffin/issues/408.
---
 crates/pypi-types/src/metadata.rs | 34 +++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/crates/pypi-types/src/metadata.rs b/crates/pypi-types/src/metadata.rs
index f63140be9..85ed37c04 100644
--- a/crates/pypi-types/src/metadata.rs
+++ b/crates/pypi-types/src/metadata.rs
@@ -211,10 +211,14 @@ impl Metadata21 {
     }
 }
 
+/// Ex) `>=7.2.0<8.0.0`
 static MISSING_COMMA: Lazy<Regex> = Lazy::new(|| Regex::new(r"(\d)([<>=~^!])").unwrap());
+/// Ex) `!=~5.0`
 static NOT_EQUAL_TILDE: Lazy<Regex> = Lazy::new(|| Regex::new(r"!=~((?:\d\.)*\d)").unwrap());
-/// e.g. `>=1.9.*`
+/// Ex) `>=1.9.*`
 static GREATER_THAN_STAR: Lazy<Regex> = Lazy::new(|| Regex::new(r">=(\d+\.\d+)\.\*").unwrap());
+/// Ex) `!=3.0*`
+static MISSING_DOT: Lazy<Regex> = Lazy::new(|| Regex::new(r"(\d\.\d)+\*").unwrap());
 
 /// Like [`Requirement`], but attempts to correct some common errors in user-provided requirements.
 #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
@@ -228,7 +232,7 @@ impl FromStr for LenientRequirement {
             Ok(requirement) => Ok(Self(requirement)),
             Err(err) => {
                 // Given `elasticsearch-dsl (>=7.2.0<8.0.0)`, rewrite to `elasticsearch-dsl (>=7.2.0,<8.0.0)`.
-                let patched = MISSING_COMMA.replace(s, r"$1,$2");
+                let patched = MISSING_COMMA.replace_all(s, r"$1,$2");
                 if patched != s {
                     if let Ok(requirement) = Requirement::from_str(&patched) {
                         warn!(
@@ -239,7 +243,7 @@ impl FromStr for LenientRequirement {
                 }
 
                 // Given `jupyter-core (!=~5.0,>=4.12)`, rewrite to `jupyter-core (!=5.0.*,>=4.12)`.
-                let patched = NOT_EQUAL_TILDE.replace(s, r"!=${1}.*");
+                let patched = NOT_EQUAL_TILDE.replace_all(s, r"!=${1}.*");
                 if patched != s {
                     if let Ok(requirement) = Requirement::from_str(&patched) {
                         warn!(
@@ -250,7 +254,7 @@ impl FromStr for LenientRequirement {
                 }
 
                 // Given `torch (>=1.9.*)`, rewrite to `torch (>=1.9)`
-                let patched = GREATER_THAN_STAR.replace(s, r">=${1}");
+                let patched = GREATER_THAN_STAR.replace_all(s, r">=${1}");
                 if patched != s {
                     if let Ok(requirement) = Requirement::from_str(&patched) {
                         warn!(
@@ -260,6 +264,17 @@ impl FromStr for LenientRequirement {
                     }
                 }
 
+                // Given `pyzmq (!=3.0*)`, rewrite to `pyzmq (!=3.0.*)`
+                let patched = MISSING_DOT.replace_all(s, r"${1}.*");
+                if patched != s {
+                    if let Ok(requirement) = Requirement::from_str(&patched) {
+                        warn!(
+                        "Inserting missing dot into invalid requirement (before: `{s}`; after: `{patched}`)",
+                    );
+                        return Ok(Self(requirement));
+                    }
+                }
+
                 Err(err)
             }
         }
@@ -313,4 +328,15 @@ mod tests {
         let expected: Requirement = Requirement::from_str("torch (>=1.9)").unwrap();
         assert_eq!(actual, expected);
     }
+
+    #[test]
+    fn missing_dot() {
+        let actual: Requirement =
+            LenientRequirement::from_str("pyzmq (>=2.7,!=3.0*,!=3.1*,!=3.2*)")
+                .unwrap()
+                .into();
+        let expected: Requirement =
+            Requirement::from_str("pyzmq (>=2.7,!=3.0.*,!=3.1.*,!=3.2.*)").unwrap();
+        assert_eq!(actual, expected);
+    }
 }