test: Add _test_simple_multiple_platforms_with_extras

python/private/pypi/extension.bzl

@@ -420,7 +420,8 @@ def _whl_repo(
             return struct(
                 repo_name = pypi_repo_name(
                     normalize_name(src.distribution),
-                    *target_platforms
+                    target_platforms = target_platforms,
+                    extras = src.extras,
                 ),
                 args = args,
                 config_setting = whl_config_setting(
@@ -450,7 +451,7 @@ def _whl_repo(
         ]
 
     return struct(
-        repo_name = whl_repo_name(src.filename, src.sha256),
+        repo_name = whl_repo_name(src.filename, src.sha256, src.extras),
         args = args,
         config_setting = whl_config_setting(
             version = python_version,

python/private/pypi/parse_requirements.bzl

@@ -133,6 +133,7 @@ def parse_requirements(
         extra_pip_args = options[target_platform]
 
         for distribution, requirement_line in reqs_:
+            req = requirement(requirement_line)
             for_whl = requirements_by_platform.setdefault(
                 normalize_name(distribution),
                 {},
@@ -146,6 +147,7 @@ def parse_requirements(
                 struct(
                     distribution = distribution,
                     srcs = index_sources(requirement_line),
+                    extras = req.extras,
                     requirement_line = requirement_line,
                     target_platforms = [],
                     extra_pip_args = extra_pip_args,
@@ -269,6 +271,7 @@ def _package_srcs(
                     distribution = name,
                     extra_pip_args = r.extra_pip_args,
                     requirement_line = req_line,
+                    extras = r.extras,
                     target_platforms = [],
                     filename = dist.filename,
                     sha256 = dist.sha256,

python/private/pypi/whl_repo_name.bzl

@@ -18,12 +18,14 @@
 load("//python/private:normalize_name.bzl", "normalize_name")
 load(":parse_whl_name.bzl", "parse_whl_name")
 
-def whl_repo_name(filename, sha256):
+def whl_repo_name(filename, sha256, extras=[]):
     """Return a valid whl_library repo name given a distribution filename.
 
     Args:
         filename: {type}`str` the filename of the distribution.
         sha256: {type}`str` the sha256 of the distribution.
+        extras: {type}`list[str]` the extras for the requirement.
+          TODO(hartikainen): Note sure if this is the right place for extras.
 
     Returns:
         a string that can be used in {obj}`whl_library`.
@@ -34,6 +36,7 @@ def whl_repo_name(filename, sha256):
         # Then the filename is basically foo-3.2.1.<ext>
         name, _, tail = filename.rpartition("-")
         parts.append(normalize_name(name))
+        parts.extend(sorted([e for e in extras if e]))
         if sha256:
             parts.append("sdist")
             version = ""
@@ -53,6 +56,7 @@ def whl_repo_name(filename, sha256):
         parts.append(python_tag)
         parts.append(abi_tag)
         parts.append(platform_tag)
+        parts.extend(sorted([e for e in extras if e]))
 
     if sha256:
         parts.append(sha256[:8])
@@ -61,19 +65,21 @@ def whl_repo_name(filename, sha256):
 
     return "_".join(parts)
 
-def pypi_repo_name(whl_name, *target_platforms):
+def pypi_repo_name(whl_name, target_platforms=[], extras=[]):
     """Return a valid whl_library given a requirement line.
 
     Args:
         whl_name: {type}`str` the whl_name to use.
-        *target_platforms: {type}`list[str]` the target platforms to use in the name.
+        target_platforms: {type}`list[str]` the target platforms to use in the name.
+        extras: {type}`list[str]` the extras for the requirement.
+          TODO(hartikainen): Note sure if this is the right place for extras.
+
 
     Returns:
         {type}`str` that can be used in {obj}`whl_library`.
     """
-    parts = [
-        normalize_name(whl_name),
-    ]
+    parts = [normalize_name(whl_name)]
+    parts.extend(sorted([e for e in extras if e]))
     parts.extend([p.partition("_")[-1] for p in target_platforms])
 
     return "_".join(parts)

tests/pypi/extension/extension_tests.bzl

@@ -407,6 +407,97 @@ new-package==0.0.1 --hash=sha256:deadb00f2
 
 _tests.append(_test_simple_multiple_python_versions)
 
+def _test_simple_multiple_extras_same_whl(env):
+    """Test that reproduces an issue where multiple extras point to same whl.
+
+    Based on https://github.com/bazel-contrib/rules_python/issues/2797#issuecomment-3143914644.
+    """
+    pypi = _parse_modules(
+        env,
+        module_ctx = _mock_mctx(
+            _mod(
+                name = "rules_python",
+                parse = [
+                    _parse(
+                        hub_name = "pypi",
+                        python_version = "3.12",
+                        download_only = True,
+                        requirements_by_platform = {
+                            "requirements.linux_arm64.txt": "linux_aarch64",
+                            "requirements.linux_x86_64.txt": "linux_x86_64",
+                        },
+                        experimental_index_url = "pypi.org",
+                    ),
+                ],
+            ),
+            read = lambda x: {
+                "requirements.linux_arm64.txt": """\
+package==0.7.0 \
+    --hash=sha256:4dd8924f171ed73a4f1a6191e2f800ae1745069989b69fabc45593d6b6504003 \
+    --hash=sha256:62833036cbaf4641d66ae94c61c0446890a91b2c0d153946583a0ebe04877a76
+""",
+                "requirements.linux_x86_64.txt": """\
+package[extra]==0.7.0 \
+    --hash=sha256:62833036cbaf4641d66ae94c61c0446890a91b2c0d153946583a0ebe04877a76
+""",
+            }[x],
+        ),
+        available_interpreters = {
+            "python_3_12_host": "unit_test_interpreter_target",
+        },
+        minor_mapping = {"3.12": "3.12.11"},
+        simpleapi_download = lambda *_, **__: {
+            "package": parse_simpleapi_html(
+                url = "https://example.com/package",
+                content = """
+<a href="package-0.7.0.tar.gz#sha256=4dd8924f171ed73a4f1a6191e2f800ae1745069989b69fabc45593d6b6504003">package-0.7.0.tar.gz</a>
+<a href="package-0.7.0-py3-none-any.whl#sha256=62833036cbaf4641d66ae94c61c0446890a91b2c0d153946583a0ebe04877a76">package-0.7.0-py3-none-any.whl</a>
+""",
+            ),
+        },
+    )
+
+    pypi.exposed_packages().contains_exactly({"pypi": ["package"]})
+    pypi.hub_whl_map().contains_exactly({"pypi": {
+        "package": {
+            "pypi_312_package_py3_none_any_62833036": [
+                whl_config_setting(
+                    target_platforms = ["cp312_linux_aarch64"],
+                    version = "3.12",
+                ),
+            ],
+            "pypi_312_package_py3_none_any_extra_62833036": [
+                whl_config_setting(
+                    target_platforms = ["cp312_linux_x86_64"],
+                    version = "3.12",
+                ),
+            ],
+        },
+    }})
+    pypi.whl_libraries().contains_exactly({
+        "pypi_312_package_py3_none_any_62833036": {
+            "dep_template": "@pypi//{name}:{target}",
+            "experimental_target_platforms": ["linux_aarch64"],
+            "filename": "package-0.7.0-py3-none-any.whl",
+            "python_interpreter_target": "unit_test_interpreter_target",
+            "requirement": "package==0.7.0",
+            "sha256": "62833036cbaf4641d66ae94c61c0446890a91b2c0d153946583a0ebe04877a76",
+            "urls": ["https://example.com/package/package-0.7.0-py3-none-any.whl"],
+        },
+        "pypi_312_package_py3_none_any_extra_62833036": {
+            "dep_template": "@pypi//{name}:{target}",
+            "experimental_target_platforms": ["linux_x86_64"],
+            "filename": "package-0.7.0-py3-none-any.whl",
+            "python_interpreter_target": "unit_test_interpreter_target",
+            "requirement": "package[extra]==0.7.0",
+            "sha256": "62833036cbaf4641d66ae94c61c0446890a91b2c0d153946583a0ebe04877a76",
+            "urls": ["https://example.com/package/package-0.7.0-py3-none-any.whl"],
+        },
+    })
+    pypi.whl_mods().contains_exactly({})
+
+_tests.append(_test_simple_multiple_extras_same_whl)
+
 def _test_simple_with_markers(env):
     pypi = _parse_modules(
         env,
@@ -1088,7 +1179,7 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     pypi.hub_whl_map().contains_exactly({
         "pypi": {
             "optimum": {
-                "pypi_315_optimum_linux_aarch64_linux_x86_64_linux_x86_64_freethreaded": [
+                "pypi_315_optimum_onnxruntime-gpu_linux_aarch64_linux_x86_64_linux_x86_64_freethreaded": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
@@ -1098,7 +1189,7 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
                         ],
                     ),
                 ],
-                "pypi_315_optimum_osx_aarch64": [
+                "pypi_315_optimum_onnxruntime_osx_aarch64": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
@@ -1111,12 +1202,12 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     })
 
     pypi.whl_libraries().contains_exactly({
-        "pypi_315_optimum_linux_aarch64_linux_x86_64_linux_x86_64_freethreaded": {
+        "pypi_315_optimum_onnxruntime-gpu_linux_aarch64_linux_x86_64_linux_x86_64_freethreaded": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime-gpu]==1.17.1",
         },
-        "pypi_315_optimum_osx_aarch64": {
+        "pypi_315_optimum_onnxruntime_osx_aarch64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime]==1.17.1",
@@ -1175,15 +1266,15 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     pypi.hub_whl_map().contains_exactly({
         "pypi": {
             "optimum": {
-                "pypi_315_optimum_mylinuxx86_64": [
+                "pypi_315_optimum_onnxruntime-gpu_mylinuxx86_64": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
                             "cp315_mylinuxx86_64",
                         ],
                     ),
                 ],
-                "pypi_315_optimum_myosxaarch64": [
+                "pypi_315_optimum_onnxruntime_myosxaarch64": [
                     whl_config_setting(
                         version = "3.15",
                         target_platforms = [
@@ -1196,12 +1287,12 @@ optimum[onnxruntime-gpu]==1.17.1 ; sys_platform == 'linux'
     })
 
     pypi.whl_libraries().contains_exactly({
-        "pypi_315_optimum_mylinuxx86_64": {
+        "pypi_315_optimum_onnxruntime-gpu_mylinuxx86_64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime-gpu]==1.17.1",
         },
-        "pypi_315_optimum_myosxaarch64": {
+        "pypi_315_optimum_onnxruntime_myosxaarch64": {
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "requirement": "optimum[onnxruntime]==1.17.1",