patch macOS wheels for <1.0.0 to specifiy CPU backend (#35)

* patch macOS wheels for `<1.0.0` to specifiy CPU backend

* add test

Author: pmeier

light_the_torch/_pip/find.py

@@ -213,6 +213,8 @@ def from_candidate_preferences(
 
 
 class PytorchCandidateEvaluator(CandidateEvaluator):
+    _MACOS_PLATFORM_PATTERN = re.compile(r"macosx_\d+_\d+_x86_64")
+
     def __init__(
         self,
         *args: Any,
@@ -245,22 +247,38 @@ def from_candidate_evaluator(
     def _sort_key(
         self, candidate: InstallationCandidate
     ) -> Tuple[cb.ComputationBackend, Version]:
-        version = Version(
-            f"{candidate.version.major}"
-            f".{candidate.version.minor}"
-            f".{candidate.version.micro}"
+        return (
+            cb.ComputationBackend.from_str(candidate.version.local),
+            candidate.version.base_version,
         )
-        computation_backend = cb.ComputationBackend.from_str(candidate.version.local)
-        return computation_backend, version
 
     def get_applicable_candidates(
         self, candidates: List[InstallationCandidate]
     ) -> List[InstallationCandidate]:
-        return [
+        applicable_candidates = [
             candidate
             for candidate in super().get_applicable_candidates(candidates)
             if candidate.version.local in self.computation_backends
         ]
+        if self._is_macos:
+            self._patch_mac_lt_1_0_0_local(applicable_candidates)
+        return applicable_candidates
+
+    @property
+    def _is_macos(self) -> bool:
+        return any(
+            self._MACOS_PLATFORM_PATTERN.match(tag.platform)
+            for tag in self._supported_tags
+        )
+
+    def _patch_mac_lt_1_0_0_local(
+        self, candidates: List[InstallationCandidate]
+    ) -> None:
+        for candidate in candidates:
+            if candidate.version.major >= 1:
+                continue
+
+            candidate.version = Version(str(candidate.version).replace("any", "cpu"))
 
 
 class PytorchLinkCollector(LinkCollector):

tests/unit/_pip/test_find.py

@@ -1,6 +1,7 @@
 import itertools
 
 import pytest
+from pip._internal.models.wheel import Wheel
 from pip._vendor.packaging.version import Version
 
 import light_the_torch as ltt
@@ -178,3 +179,19 @@ def test_find_links_channel_smoke(channel):
     assert ltt.find_links(
         ["torch"], computation_backends={cb.CPUBackend()}, channel=channel
     )
+
+
+@pytest.mark.parametrize("python_version", PYTHON_VERSIONS)
+def test_patch_mac_local_specifier_lt_1_0_0(
+    patch_extract_dists, patch_run, python_version
+):
+    # See https://github.com/pmeier/light-the-torch/issues/34
+    dists = ["torch"]
+    patch_extract_dists(return_value=dists)
+
+    links = ltt.find_links(
+        dists, python_version=python_version, platform="macosx_10_9_x86_64"
+    )
+    version = Version(Wheel(links[0]).version)
+
+    assert version >= Version("1.0.0")