From 0a4c0fe7ccd72a9ff78ee2ee5d166ca9c4f194d0 Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Fri, 8 Nov 2024 21:09:53 +0100
Subject: [PATCH] Add option to specify additional markers for wheel build
 requirements (#129949)

---
 script/gen_requirements_all.py            | 35 +++++++++++++++++++----
 tests/script/test_gen_requirements_all.py | 26 +++++++++++++++++
 2 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/script/gen_requirements_all.py b/script/gen_requirements_all.py
index 4a340863240..02dad3aef3f 100755
--- a/script/gen_requirements_all.py
+++ b/script/gen_requirements_all.py
@@ -58,8 +58,16 @@ INCLUDED_REQUIREMENTS_WHEELS = {
 # will be included in requirements_all_{action}.txt
 
 OVERRIDDEN_REQUIREMENTS_ACTIONS = {
-    "pytest": {"exclude": set(), "include": {"python-gammu"}},
-    "wheels_aarch64": {"exclude": set(), "include": INCLUDED_REQUIREMENTS_WHEELS},
+    "pytest": {
+        "exclude": set(),
+        "include": {"python-gammu"},
+        "markers": {},
+    },
+    "wheels_aarch64": {
+        "exclude": set(),
+        "include": INCLUDED_REQUIREMENTS_WHEELS,
+        "markers": {},
+    },
     # Pandas has issues building on armhf, it is expected they
     # will drop the platform in the near future (they consider it
     # "flimsy" on 386). The following packages depend on pandas,
@@ -67,10 +75,23 @@ OVERRIDDEN_REQUIREMENTS_ACTIONS = {
     "wheels_armhf": {
         "exclude": {"env-canada", "noaa-coops", "pyezviz", "pykrakenapi"},
         "include": INCLUDED_REQUIREMENTS_WHEELS,
+        "markers": {},
+    },
+    "wheels_armv7": {
+        "exclude": set(),
+        "include": INCLUDED_REQUIREMENTS_WHEELS,
+        "markers": {},
+    },
+    "wheels_amd64": {
+        "exclude": set(),
+        "include": INCLUDED_REQUIREMENTS_WHEELS,
+        "markers": {},
+    },
+    "wheels_i386": {
+        "exclude": set(),
+        "include": INCLUDED_REQUIREMENTS_WHEELS,
+        "markers": {},
     },
-    "wheels_armv7": {"exclude": set(), "include": INCLUDED_REQUIREMENTS_WHEELS},
-    "wheels_amd64": {"exclude": set(), "include": INCLUDED_REQUIREMENTS_WHEELS},
-    "wheels_i386": {"exclude": set(), "include": INCLUDED_REQUIREMENTS_WHEELS},
 }
 
 IGNORE_PIN = ("colorlog>2.1,<3", "urllib3")
@@ -311,6 +332,10 @@ def process_action_requirement(req: str, action: str) -> str:
         return req
     if normalized_package_name in EXCLUDED_REQUIREMENTS_ALL:
         return f"# {req}"
+    if markers := OVERRIDDEN_REQUIREMENTS_ACTIONS[action]["markers"].get(
+        normalized_package_name, None
+    ):
+        return f"{req};{markers}"
     return req
 
 
diff --git a/tests/script/test_gen_requirements_all.py b/tests/script/test_gen_requirements_all.py
index 793b3de63c5..519a5c21855 100644
--- a/tests/script/test_gen_requirements_all.py
+++ b/tests/script/test_gen_requirements_all.py
@@ -1,5 +1,7 @@
 """Tests for the gen_requirements_all script."""
 
+from unittest.mock import patch
+
 from script import gen_requirements_all
 
 
@@ -23,3 +25,27 @@ def test_include_overrides_subsets() -> None:
     for overrides in gen_requirements_all.OVERRIDDEN_REQUIREMENTS_ACTIONS.values():
         for req in overrides["include"]:
             assert req in gen_requirements_all.EXCLUDED_REQUIREMENTS_ALL
+
+
+def test_requirement_override_markers() -> None:
+    """Test override markers are applied to the correct requirements."""
+    data = {
+        "pytest": {
+            "exclude": set(),
+            "include": set(),
+            "markers": {"env-canada": "python_version<'3.13'"},
+        }
+    }
+    with patch.dict(
+        gen_requirements_all.OVERRIDDEN_REQUIREMENTS_ACTIONS, data, clear=True
+    ):
+        assert (
+            gen_requirements_all.process_action_requirement(
+                "env-canada==0.7.2", "pytest"
+            )
+            == "env-canada==0.7.2;python_version<'3.13'"
+        )
+        assert (
+            gen_requirements_all.process_action_requirement("other==1.0", "pytest")
+            == "other==1.0"
+        )