1 files changed, 30 insertions, 6 deletions
diff --git a/pkgs/development/python-modules/mlrose/default.nix b/pkgs/development/python-modules/mlrose/default.nix
index 40e3b6ce0953..2187f2726801 100644
--- a/pkgs/development/python-modules/mlrose/default.nix
+++ b/pkgs/development/python-modules/mlrose/default.nix
@@ -1,21 +1,45 @@
-{ lib, isPy27, buildPythonPackage, fetchPypi, scikitlearn }:
+{ lib
+, isPy27
+, buildPythonPackage
+, fetchFromGitHub
+, fetchpatch
+, scikit-learn
+, pytestCheckHook
+, pytest-randomly
+}:
 
 buildPythonPackage rec {
   pname = "mlrose";
   version = "1.3.0";
   disabled = isPy27;
 
-  src = fetchPypi {
-    inherit pname version;
-    sha256 = "cec83253bf6da67a7fb32b2c9ae13e9dbc6cfbcaae2aa3107993e69e9788f15e";
+  src = fetchFromGitHub {
+    owner = "gkhayes";
+    repo = "mlrose";
+    rev = "v${version}";
+    sha256 = "1dn43k3rcypj58ymcj849b37w66jz7fphw8842v6mlbij3x0rxfl";
   };
 
-  propagatedBuildInputs = [ scikitlearn ];
+  patches = [
+    # Fixes compatibility with scikit-learn 0.24.1
+    (fetchpatch {
+      url = "https://github.com/gkhayes/mlrose/pull/55/commits/19caf8616fc194402678aa67917db334ad02852a.patch";
+      sha256 = "1nivz3bn21nd21bxbcl16a6jmy7y5j8ilz90cjmd0xq4v7flsahf";
+    })
+  ];
+
+  propagatedBuildInputs = [ scikit-learn ];
+  checkInputs = [ pytest-randomly pytestCheckHook ];
 
   postPatch = ''
-    sed -i 's,sklearn,scikit-learn,g' setup.py
+    substituteInPlace setup.py --replace sklearn scikit-learn
   '';
 
+  pythonImportsCheck = [ "mlrose" ];
+
+  # Fix random seed during tests
+  pytestFlagsArray = [ "--randomly-seed 0" ];
+
   meta = with lib; {
     description = "Machine Learning, Randomized Optimization and SEarch";
     homepage    = "https://github.com/gkhayes/mlrose";