From 6f1485cc1331e61b8fa92c7d5dafe44f5ede8ac0 Mon Sep 17 00:00:00 2001
From: Christian Duerr <chris.durr@phylum.io>
Date: Wed, 11 Oct 2023 18:00:28 +0200
Subject: [PATCH] Allow for equals in param

---
 lockfile/src/parsers/pypi.rs           |  8 ++++++--
 lockfile/src/python.rs                 | 10 +++++++++-
 tests/fixtures/requirements-locked.txt |  4 +++-
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/lockfile/src/parsers/pypi.rs b/lockfile/src/parsers/pypi.rs
index c6b844a25..e80aabbaf 100644
--- a/lockfile/src/parsers/pypi.rs
+++ b/lockfile/src/parsers/pypi.rs
@@ -57,12 +57,16 @@ fn line<'a>(input: &'a str, registry: &mut Option<&'a str>) -> IResult<&'a str,
     //
     // Since `ThirdPartyVersion` only allows a single registry, we only record the
     // primary one.
-    if let Some(index_url) = line.strip_prefix("--index-url ").or_else(|| line.strip_prefix("-i "))
+    if let Some(mut index_url) =
+        line.strip_prefix("--index-url").or_else(|| line.strip_prefix("-i"))
     {
+        // Strip optional `=` from `--index-url=https://...`
+        index_url = index_url.trim().strip_prefix('=').unwrap_or(index_url).trim();
+
         *registry = Some(index_url.trim());
         line = "";
     }
-    if line.starts_with("--extra-index-url ") {
+    if line.starts_with("--extra-index-url") {
         line = "";
     }
 
diff --git a/lockfile/src/python.rs b/lockfile/src/python.rs
index ea55da30f..09b84b4f0 100644
--- a/lockfile/src/python.rs
+++ b/lockfile/src/python.rs
@@ -238,7 +238,7 @@ mod tests {
         let pkgs = PyRequirements
             .parse(include_str!("../../tests/fixtures/requirements-locked.txt"))
             .unwrap();
-        assert_eq!(pkgs.len(), 13);
+        assert_eq!(pkgs.len(), 14);
 
         let expected_pkgs = [
             Package {
@@ -301,6 +301,14 @@ mod tests {
                 version: PackageVersion::Path(Some("/tmp/editable".into())),
                 package_type: PackageType::PyPi,
             },
+            Package {
+                name: "other-registry-a".into(),
+                version: PackageVersion::ThirdParty(ThirdPartyVersion {
+                    registry: "https://mirror1.phylum.io/simple/".into(),
+                    version: "3.2.1".into(),
+                }),
+                package_type: PackageType::PyPi,
+            },
             Package {
                 name: "other-registry".into(),
                 version: PackageVersion::ThirdParty(ThirdPartyVersion {
diff --git a/tests/fixtures/requirements-locked.txt b/tests/fixtures/requirements-locked.txt
index c1515d4ee..797df1d49 100644
--- a/tests/fixtures/requirements-locked.txt
+++ b/tests/fixtures/requirements-locked.txt
@@ -34,6 +34,8 @@ tomli @ https://files.pythonhosted.org/packages/97/75/10a9ebee3fd790d20926a90a25
 -e /tmp/editable ; python_version >= "3.7" and python_version < "3.12"
 
 --index-url https://mirror1.phylum.io/simple/
--i https://mirror2.phylum.io/simple/
+other-registry-a==3.2.1
+-i=https://mirror2.phylum.io/simple/
 --extra-index-url https://mirror3.phylum.io/simple/
+--extra-index-url=https://mirror3.phylum.io/simple/
 other-registry==1.2.3