From 170dd558dffd2bb134e3001a93631fc71fe13f3c Mon Sep 17 00:00:00 2001
From: Nikolay Amiantov <ab@fmap.me>
Date: Sat, 13 Jul 2019 22:37:01 +0300
Subject: [PATCH] python.pkgs.tensorflow: cleanup binary build

---
 .../python-modules/tensorflow/bin.nix         | 25 +++++++++++--------
 ...{tf1.14.0-hashes.nix => binary-hashes.nix} |  1 +
 .../python-modules/tensorflow/prefetcher.sh   |  5 +++-
 3 files changed, 19 insertions(+), 12 deletions(-)
 rename pkgs/development/python-modules/tensorflow/{tf1.14.0-hashes.nix => binary-hashes.nix} (98%)
 mode change 100644 => 100755 pkgs/development/python-modules/tensorflow/prefetcher.sh

diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix
index 817ba09dc9d2..a562ba8f2fad 100644
--- a/pkgs/development/python-modules/tensorflow/bin.nix
+++ b/pkgs/development/python-modules/tensorflow/bin.nix
@@ -35,15 +35,19 @@
 assert cudaSupport -> cudatoolkit != null
                    && cudnn != null
                    && nvidia_x11 != null;
+
+# unsupported combination
+assert ! (stdenv.isDarwin && cudaSupport);
+
 let
-  cudatoolkit_joined = symlinkJoin {
-    name = "unsplit_cudatoolkit";
-    paths = [ cudatoolkit.out
-              cudatoolkit.lib ];};
+  packages = import ./binary-hashes.nix;
+
+  variant = if cudaSupport then "-gpu" else "";
+  pname = "tensorflow${variant}";
 
 in buildPythonPackage rec {
-  pname = "tensorflow";
-  version = "1.14.0";
+  inherit pname;
+  inherit (packages) version;
   format = "wheel";
 
   src = let
@@ -52,8 +56,7 @@ in buildPythonPackage rec {
     platform = if stdenv.isDarwin then "mac" else "linux";
     unit = if cudaSupport then "gpu" else "cpu";
     key = "${platform}_py_${pyver}_${unit}";
-    dls = import (./. + "/tf${version}-hashes.nix");
-  in fetchurl dls.${key};
+  in fetchurl packages.${key};
 
   propagatedBuildInputs = [
     protobuf
@@ -86,9 +89,9 @@ in buildPythonPackage rec {
   # patchelf --shrink-rpath will remove the cuda libraries.
   postFixup = let
     rpath = stdenv.lib.makeLibraryPath
-      ([ stdenv.cc.cc.lib zlib ] ++ lib.optionals cudaSupport [ cudatoolkit_joined cudnn nvidia_x11 ]);
+      ([ stdenv.cc.cc.lib zlib ] ++ lib.optionals cudaSupport [ cudatoolkit.out cudatoolkit.lib cudnn nvidia_x11 ]);
   in
-  lib.optionalString (stdenv.isLinux) ''
+  lib.optionalString stdenv.isLinux ''
     rrPath="$out/${python.sitePackages}/tensorflow/:$out/${python.sitePackages}/tensorflow/contrib/tensor_forest/:${rpath}"
     internalLibPath="$out/${python.sitePackages}/tensorflow/python/_pywrap_tensorflow_internal.so"
     find $out \( -name '*.so' -or -name '*.so.*' \) -exec patchelf --set-rpath "$rrPath" {} \;
@@ -100,7 +103,7 @@ in buildPythonPackage rec {
     homepage = http://tensorflow.org;
     license = licenses.asl20;
     maintainers = with maintainers; [ jyp abbradar ];
-    platforms = with platforms; linux ++ lib.optionals (!cudaSupport) darwin;
+    platforms = [ "x86_64-linux" "x86_64-darwin" ];
     # Python 2.7 build uses different string encoding.
     # See https://github.com/NixOS/nixpkgs/pull/37044#issuecomment-373452253
     broken = stdenv.isDarwin && !isPy3k;
diff --git a/pkgs/development/python-modules/tensorflow/tf1.14.0-hashes.nix b/pkgs/development/python-modules/tensorflow/binary-hashes.nix
similarity index 98%
rename from pkgs/development/python-modules/tensorflow/tf1.14.0-hashes.nix
rename to pkgs/development/python-modules/tensorflow/binary-hashes.nix
index f6b367aea383..fa4809dc3f19 100644
--- a/pkgs/development/python-modules/tensorflow/tf1.14.0-hashes.nix
+++ b/pkgs/development/python-modules/tensorflow/binary-hashes.nix
@@ -1,4 +1,5 @@
 {
+version = "1.14.0";
 linux_py_27_cpu = {
   url = "https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.14.0-cp27-none-linux_x86_64.whl";
   sha256 = "0yywdrfk97dh1bxhibspg0raz70fx9lcczj6xlimqy4xb60clx7k";
diff --git a/pkgs/development/python-modules/tensorflow/prefetcher.sh b/pkgs/development/python-modules/tensorflow/prefetcher.sh
old mode 100644
new mode 100755
index 3059340d127d..d590fb0f1732
--- a/pkgs/development/python-modules/tensorflow/prefetcher.sh
+++ b/pkgs/development/python-modules/tensorflow/prefetcher.sh
@@ -1,7 +1,10 @@
+#!/usr/bin/env bash
+
 version=1.14.0
-hashfile=tf${version}-hashes.nix
+hashfile=binary-hashes.nix
 rm -f $hashfile
 echo "{" >> $hashfile
+echo "version = \"$version\";" >> $hashfile
 for sys in "linux" "mac"; do
     for tfpref in "cpu/tensorflow" "gpu/tensorflow_gpu"; do
         for pykind in "py2-none-any" "py3-none-any" "cp27-none-linux_x86_64" "cp35-cp35m-linux_x86_64" "cp36-cp36m-linux_x86_64" "cp37-cp37m-linux_x86_64"; do