From e677462813f7bfcdf4db5f848f641740c1e573d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sandro=20J=C3=A4ckel?= <sandro.jaeckel@gmail.com>
Date: Mon, 22 Feb 2021 08:39:00 +0100
Subject: [PATCH] pythonPackages.sagemaker: add pandas for analytical feature

---
 .../python-modules/sagemaker/default.nix          | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/pkgs/development/python-modules/sagemaker/default.nix b/pkgs/development/python-modules/sagemaker/default.nix
index 6b73642aea65..5df445f299ba 100644
--- a/pkgs/development/python-modules/sagemaker/default.nix
+++ b/pkgs/development/python-modules/sagemaker/default.nix
@@ -9,6 +9,7 @@
 , protobuf
 , protobuf3-to-dict
 , smdebug-rulesconfig
+, pandas
 }:
 
 buildPythonPackage rec {
@@ -20,9 +21,10 @@ buildPythonPackage rec {
     sha256 = "sha256-xQ1nt8FcjuoilzM5PbU8KHgirPyj9us+ykyjfgEqZhg=";
   };
 
-  doCheck = false;
-
-  pythonImportsCheck = [ "sagemaker" ];
+  pythonImportsCheck = [
+    "sagemaker"
+    "sagemaker.lineage.visualizer"
+  ];
 
   propagatedBuildInputs = [
     attrs
@@ -33,8 +35,15 @@ buildPythonPackage rec {
     protobuf
     protobuf3-to-dict
     smdebug-rulesconfig
+    pandas
   ];
 
+  doCheck = false;
+
+  postFixup = ''
+    [ "$($out/bin/sagemaker-upgrade-v2 --help 2>&1 | grep -cim1 'pandas failed to import')" -eq "0" ]
+  '';
+
   meta = with lib; {
     description = "Library for training and deploying machine learning models on Amazon SageMaker";
     homepage = "https://github.com/aws/sagemaker-python-sdk/";