From 3a1eb66fbc98089fa24f13243513af40a633482a Mon Sep 17 00:00:00 2001
From: Tobias Pflug <tobias.pflug@gmail.com>
Date: Tue, 17 Dec 2019 18:13:11 +0100
Subject: [PATCH] Improve PEP425 support

Make sure that whl files using notation such as "py2|py3|py2.py3"
and/or files using "any" in the `abi` tag are accepted as well.
---
 pep425.nix               | 27 +++++++-----
 tests/pep425/default.nix | 89 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+), 11 deletions(-)

diff --git a/pep425.nix b/pep425.nix
index 052c5a5..b2e1120 100644
--- a/pep425.nix
+++ b/pep425.nix
@@ -33,12 +33,6 @@ let
         platform = p;
       };
 
-  #
-  # Renders a wheel attribute set into a string
-  #
-  # e.g (toFile (toWheelAttrs x)) == x
-  toFileName = x: "${x.pkgName}-${x.pkgVer}-${x.pyVer}-${x.abi}-${x.platform}.whl";
-
   #
   # Builds list of acceptable osx wheel files
   #
@@ -53,6 +47,14 @@ let
       then []
       else (builtins.filter (x: hasInfix v x.file) candidates) ++ (findBestMatches vs candidates);
 
+  # pyver = "cpXX"
+  # x     = "cpXX" | "py2" | "py3" | "py2.py3"
+  isPyVersionCompatible = pyver: x:
+    let
+      normalize = y: ''cp${lib.strings.removePrefix "cp" (lib.strings.removePrefix "py" y)}'';
+      isCompat = p: x: lib.strings.hasPrefix (normalize x) p;
+    in
+      lib.lists.any (isCompat pyver) (lib.strings.splitString "." x);
 
   #
   # Selects the best matching wheel file from a list of files
@@ -61,15 +63,18 @@ let
     let
       filesWithoutSources = (builtins.filter (x: hasSuffix ".whl" x.file) files);
 
-      withPython = pyver: abi: x: x.pyVer == pyver && x.abi == abi;
+      isPyAbiCompatible = pyabi: x: x == "none" || pyabi == x;
+
+      withPython = ver: abi: x: (isPyVersionCompatible ver x.pyVer) && (isPyAbiCompatible abi x.abi);
 
       withPlatform = if isLinux
       then (
         x: x.platform == "manylinux1_${stdenv.platform.kernelArch}"
         || x.platform == "manylinux2010_${stdenv.platform.kernelArch}"
         || x.platform == "manylinux2014_${stdenv.platform.kernelArch}"
+        || x.platform == "any"
       )
-      else (x: hasInfix "macosx" x.platform);
+      else (x: hasInfix "macosx" x.platform || x.platform == "any");
 
       filterWheel = x:
         let
@@ -81,8 +86,8 @@ let
 
       choose = files:
         let
-          osxMatches = [ "10_12" "10_11" "10_10" "10_9" ];
-          linuxMatches = [ "manylinux1_" "manylinux2010_" "manylinux2014_" ];
+          osxMatches = [ "10_12" "10_11" "10_10" "10_9" "any" ];
+          linuxMatches = [ "manylinux1_" "manylinux2010_" "manylinux2014_" "any" ];
           chooseLinux = x: lib.singleton (builtins.head (findBestMatches linuxMatches x));
           chooseOSX = x: lib.singleton (builtins.head (findBestMatches osxMatches x));
         in
@@ -97,5 +102,5 @@ let
 
 in
 {
-  inherit selectWheel;
+  inherit selectWheel toWheelAttrs isPyVersionCompatible;
 }
diff --git a/tests/pep425/default.nix b/tests/pep425/default.nix
index 8166adf..623cb57 100644
--- a/tests/pep425/default.nix
+++ b/tests/pep425/default.nix
@@ -2,6 +2,10 @@
 
 lib.debug.runTests {
 
+  #
+  # selectWheel
+  #
+
   testLinuxSimple =
     let
       cs = [
@@ -173,4 +177,89 @@ lib.debug.runTests {
         expr = pep425Python37.selectWheel cs;
         expected = [ { file = "msgpack-0.6.2-cp37-cp37m-manylinux1_x86_64.whl"; } ];
       };
+
+  testNonManyLinuxWheels =
+    let
+      cs = [
+        { file = "tensorboard-1.14.0-py2-none-any.whl"; }
+        { file = "tensorboard-1.14.0-py3-none-any.whl"; }
+      ];
+    in
+      {
+        expr = pep425Python37.selectWheel cs;
+        expected = [ { file = "tensorboard-1.14.0-py3-none-any.whl"; } ];
+      };
+
+  testPy2Py3Wheels =
+    let
+      cs = [
+        { file = "tensorboard-1.14.0-py2.py3-none-any.whl"; }
+      ];
+    in
+      {
+        expr = pep425Python37.selectWheel cs;
+        expected = [ { file = "tensorboard-1.14.0-py2.py3-none-any.whl"; } ];
+      };
+
+  #
+  # toWheelAttrs
+  #
+
+  testToWheelAttrs =
+    let
+      name = "msgpack-0.6.2-cp27-cp27m-manylinux1_i686.whl";
+    in
+      {
+        expr = pep425.toWheelAttrs name;
+        expected = {
+          pkgName = "msgpack";
+          pkgVer = "0.6.2";
+          pyVer = "cp27";
+          abi = "cp27m";
+          platform = "manylinux1_i686";
+        };
+      };
+
+  testToWheelAttrsAny =
+    let
+      name = "tensorboard-1.14.0-py3-none-any.whl";
+    in
+      {
+        expr = pep425.toWheelAttrs name;
+        expected = {
+          pkgName = "tensorboard";
+          pkgVer = "1.14.0";
+          pyVer = "py3";
+          abi = "none";
+          platform = "any";
+        };
+      };
+
+  #
+  # isPyVersionCompatible
+  #
+
+  tesPyVersionCompatible =
+    let
+      f = pep425.isPyVersionCompatible;
+    in
+      {
+        expr = [
+          (f "cp27" "cp27")
+          (f "cp27" "cp37")
+          (f "cp27" "py2")
+          (f "cp27" "py3")
+          (f "cp27" "py2.py3")
+          (f "cp37" "py2.py3")
+        ];
+
+        expected = [
+          true
+          false
+          true
+          false
+          true
+        ];
+      };
+
 }