Allow initializing WheelPolicies with a specific libc & architecture

Author: mayeut

src/auditwheel/policy/__init__.py

@@ -17,6 +17,7 @@
 
 _HERE = Path(__file__).parent
 LIBPYTHON_RE = re.compile(r"^libpython\d+\.\d+m?.so(.\d)*$")
+_MUSL_POLICY_RE = re.compile(r"^musllinux_\d+_\d+$")
 
 logger = logging.getLogger(__name__)
 
@@ -30,14 +31,30 @@
 
 
 class WheelPolicies:
-    def __init__(self) -> None:
-        libc_variant = get_libc()
-        policies_path = _POLICY_JSON_MAP[libc_variant]
-        policies = json.loads(policies_path.read_text())
+    def __init__(
+        self,
+        *,
+        libc: Libc | None = None,
+        musl_policy: str | None = None,
+        arch: str | None = None,
+    ) -> None:
+        if libc is None:
+            libc = get_libc() if musl_policy is None else Libc.MUSL
+        if libc != Libc.MUSL and musl_policy is not None:
+            raise ValueError(f"'musl_policy' shall be None for libc {libc.name}")
+        if libc == Libc.MUSL:
+            if musl_policy is None:
+                musl_version = get_musl_version(find_musl_libc())
+                musl_policy = f"musllinux_{musl_version.major}_{musl_version.minor}"
+            elif _MUSL_POLICY_RE.match(musl_policy) is None:
+                raise ValueError(f"Invalid 'musl_policy': '{musl_policy}'")
+        if arch is None:
+            arch = get_arch_name()
+        policies = json.loads(_POLICY_JSON_MAP[libc].read_text())
         self._policies = []
-        self._musl_policy = _get_musl_policy()
-        self._arch_name = get_arch_name()
-        self._libc_variant = get_libc()
+        self._arch_name = arch
+        self._libc_variant = libc
+        self._musl_policy = musl_policy
 
         _validate_pep600_compliance(policies)
         for policy in policies:
@@ -59,7 +76,7 @@ def __init__(self) -> None:
                     alias + "_" + self._arch_name for alias in policy["aliases"]
                 ]
                 policy["lib_whitelist"] = _fixup_musl_libc_soname(
-                    policy["lib_whitelist"]
+                    libc, arch, policy["lib_whitelist"]
                 )
                 self._policies.append(policy)
 
@@ -219,10 +236,6 @@ def get_arch_name() -> str:
     return machine
 
 
-_ARCH_NAME = get_arch_name()
-_LIBC = get_libc()
-
-
 def _validate_pep600_compliance(policies) -> None:
     symbol_versions: dict[str, dict[str, set[str]]] = {}
     lib_whitelist: set[str] = set()
@@ -253,15 +266,8 @@ def _validate_pep600_compliance(policies) -> None:
             symbol_versions[arch] = symbol_versions_arch
 
 
-def _get_musl_policy():
-    if _LIBC != Libc.MUSL:
-        return None
-    musl_version = get_musl_version(find_musl_libc())
-    return f"musllinux_{musl_version.major}_{musl_version.minor}"
-
-
-def _fixup_musl_libc_soname(whitelist):
-    if _LIBC != Libc.MUSL:
+def _fixup_musl_libc_soname(libc: Libc, arch: str, whitelist):
+    if libc != Libc.MUSL:
         return whitelist
     soname_map = {
         "libc.so": {
@@ -276,7 +282,7 @@ def _fixup_musl_libc_soname(whitelist):
     new_whitelist = []
     for soname in whitelist:
         if soname in soname_map:
-            new_soname = soname_map[soname][_ARCH_NAME]
+            new_soname = soname_map[soname][arch]
             logger.debug(f"Replacing whitelisted '{soname}' by '{new_soname}'")
             new_whitelist.append(new_soname)
         else:

tests/integration/test_bundled_wheels.py

@@ -13,13 +13,13 @@
 import pytest
 
 from auditwheel import main_repair
+from auditwheel.libc import Libc
 from auditwheel.policy import WheelPolicies
 from auditwheel.wheel_abi import analyze_wheel_abi
 
 HERE = Path(__file__).parent.resolve()
 
 
-@pytest.mark.skipif(platform.machine() != "x86_64", reason="only supported on x86_64")
 @pytest.mark.parametrize(
     "file, external_libs",
     [
@@ -28,14 +28,13 @@
     ],
 )
 def test_analyze_wheel_abi(file, external_libs):
-    wheel_policies = WheelPolicies()
+    wheel_policies = WheelPolicies(libc=Libc.GLIBC, arch="x86_64")
     winfo = analyze_wheel_abi(wheel_policies, str(HERE / file))
     assert set(winfo.external_refs["manylinux_2_5_x86_64"]["libs"]) == external_libs
 
 
-@pytest.mark.skipif(platform.machine() != "x86_64", reason="only supported on x86_64")
 def test_analyze_wheel_abi_pyfpe():
-    wheel_policies = WheelPolicies()
+    wheel_policies = WheelPolicies(libc=Libc.GLIBC, arch="x86_64")
     winfo = analyze_wheel_abi(
         wheel_policies, str(HERE / "fpewheel-0.0.0-cp35-cp35m-linux_x86_64.whl")
     )

tests/unit/test_policy.py

@@ -1,17 +1,37 @@
 from __future__ import annotations
 
+import re
+from contextlib import nullcontext as does_not_raise
 from unittest.mock import patch
 
 import pytest
 
+from auditwheel.error import InvalidLibc
+from auditwheel.libc import Libc
 from auditwheel.policy import (
     WheelPolicies,
     _validate_pep600_compliance,
     get_arch_name,
+    get_libc,
     get_replace_platforms,
 )
 
 
+def ids(x):
+    if isinstance(x, Libc):
+        return x.name
+    if isinstance(x, does_not_raise):
+        return "NoError"
+    if hasattr(x, "expected_exception"):
+        return x.expected_exception
+
+
+def raises(exception, match=None, escape=True):
+    if escape and match is not None:
+        match = re.escape(match)
+    return pytest.raises(exception, match=match)
+
+
 @patch("auditwheel.policy._platform_module.machine")
 @patch("auditwheel.policy.bits", 32)
 @pytest.mark.parametrize(
@@ -233,3 +253,43 @@ def test_filter_libs(self):
         # Assert that each policy only has the unfiltered libs.
         for policy in full_external_refs:
             assert set(full_external_refs[policy]["libs"]) == set(unfiltered_libs)
+
+
+@pytest.mark.parametrize(
+    "libc,musl_policy,arch,exception",
+    [
+        # valid
+        (None, None, None, does_not_raise()),
+        (Libc.GLIBC, None, None, does_not_raise()),
+        (Libc.MUSL, "musllinux_1_1", None, does_not_raise()),
+        (None, "musllinux_1_1", None, does_not_raise()),
+        (None, None, "aarch64", does_not_raise()),
+        # invalid
+        (
+            Libc.GLIBC,
+            "musllinux_1_1",
+            None,
+            raises(ValueError, "'musl_policy' shall be None"),
+        ),
+        (Libc.MUSL, "manylinux_1_1", None, raises(ValueError, "Invalid 'musl_policy'")),
+        (Libc.MUSL, "musllinux_5_1", None, raises(AssertionError)),
+        (Libc.MUSL, "musllinux_1_1", "foo", raises(AssertionError)),
+        # platform dependant
+        (
+            Libc.MUSL,
+            None,
+            None,
+            does_not_raise() if get_libc() == Libc.MUSL else raises(InvalidLibc),
+        ),
+    ],
+    ids=ids,
+)
+def test_wheel_policies_args(libc, musl_policy, arch, exception):
+    with exception:
+        wheel_policies = WheelPolicies(libc=libc, musl_policy=musl_policy, arch=arch)
+        if libc is not None:
+            assert wheel_policies._libc_variant == libc
+        if musl_policy is not None:
+            assert wheel_policies._musl_policy == musl_policy
+        if arch is not None:
+            assert wheel_policies._arch_name == arch