feat(pypi) Handle local version in requirement parsing

Author: ewianda

python/private/pypi/parse_requirements.bzl

@@ -32,6 +32,15 @@ load(":index_sources.bzl", "index_sources")
 load(":parse_requirements_txt.bzl", "parse_requirements_txt")
 load(":whl_target_platforms.bzl", "select_whls")
 
+def extract_version(entry):
+    """Extract the version part from the requirement string."""
+    version_start = entry.find("==")
+    if version_start != -1:
+        # Extract everything after '==' until the next space or end of the string
+        version = entry[version_start + 2:].split(" ")[0]
+        return version
+    return None
+
 def parse_requirements(
         ctx,
         *,
@@ -92,7 +101,7 @@ def parse_requirements(
         # are returned as just the base package name. e.g., `foo[bar]` results
         # in an entry like `("foo", "foo[bar] == 1.0 ...")`.
         requirements_dict = {
-            normalize_name(entry[0]): entry
+            (entry[0], extract_version(entry[1])): entry
             for entry in sorted(
                 parse_result.requirements,
                 # Get the longest match and fallback to original WORKSPACE sorting,

tests/pypi/parse_requirements/parse_requirements_tests.bzl

@@ -66,6 +66,10 @@ foo==0.0.3 --hash=sha256:deadbaaf
         "requirements_windows": """\
 foo[extra]==0.0.2 --hash=sha256:deadbeef
 bar==0.0.1 --hash=sha256:deadb00f
+""",
+        "requirements_different_package_version": """\
+foo==0.0.1+local --hash=sha256:deadbeef
+foo==0.0.1 --hash=sha256:deadb00f
 """,
     }
 
@@ -382,6 +386,49 @@ def _test_env_marker_resolution(env):
 
 _tests.append(_test_env_marker_resolution)
 
+#  <list [struct(distribution = "foo", extra_pip_args = [], is_exposed = True, requirement_line = "foo==0.0.1 --hash=sha256:deadb00f", sdist = None, srcs = struct(requirement = "foo==0.0.1", shas = ["deadb00f"], version = "0.0.1"), target_platforms = ["linux_x86_64"], whls = []), struct(distribution = "foo", extra_pip_args = [], is_exposed = True, requirement_line = "foo==0.0.1+local --hash=sha256:deadbeef", sdist = None, srcs = struct(requirement = "foo==0.0.1+local", shas = ["deadbeef"], version = "0.0.1+local"), target_platforms = ["linux_x86_64"], whls = [])]>
+def _test_different_package_version(env):
+    got = parse_requirements(
+        ctx = _mock_ctx(),
+        requirements_by_platform = {
+            "requirements_different_package_version": ["linux_x86_64"],
+        },
+    )
+    env.expect.that_dict(got).contains_exactly({
+        "foo": [
+            struct(
+                distribution = "foo",
+                extra_pip_args = [],
+                requirement_line = "foo==0.0.1 --hash=sha256:deadb00f",
+                srcs = struct(
+                    requirement = "foo==0.0.1",
+                    shas = ["deadb00f"],
+                    version = "0.0.1",
+                ),
+                target_platforms = ["linux_x86_64"],
+                whls = [],
+                sdist = None,
+                is_exposed = True,
+            ),
+            struct(
+                distribution = "foo",
+                extra_pip_args = [],
+                requirement_line = "foo==0.0.1+local --hash=sha256:deadbeef",
+                srcs = struct(
+                    requirement = "foo==0.0.1+local",
+                    shas = ["deadbeef"],
+                    version = "0.0.1+local",
+                ),
+                target_platforms = ["linux_x86_64"],
+                whls = [],
+                sdist = None,
+                is_exposed = True,
+            ),
+        ],
+    })
+
+_tests.append(_test_different_package_version)
+
 def parse_requirements_test_suite(name):
     """Create the test suite.