python3Packages.skrl: cleanup, fix

Author: GaetanLepage

pkgs/development/python-modules/skrl/default.nix

@@ -2,68 +2,80 @@
   lib,
   buildPythonPackage,
   fetchFromGitHub,
+
+  # build-system
   setuptools,
+
+  # dependencies
   gym,
   gymnasium,
-  torch,
+  packaging,
   tensorboard,
+  torch,
   tqdm,
   wandb,
-  packaging,
+
+  # tests
+  flax,
+  jax,
+  optax,
+  pettingzoo,
+  pygame,
+  pymunk,
   pytestCheckHook,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "skrl";
   version = "1.4.3";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "Toni-SM";
     repo = "skrl";
-    tag = version;
+    tag = finalAttrs.version;
     hash = "sha256-5lkoYAmMIWqK3+E3WxXMWS9zal2DhZkfl30EkrHKpdI=";
   };
 
-  nativeBuildInputs = [ setuptools ];
+  build-system = [ setuptools ];
 
-  propagatedBuildInputs = [
+  dependencies = [
     gym
     gymnasium
-    torch
+    packaging
     tensorboard
+    torch
     tqdm
     wandb
-    packaging
   ];
 
-  nativeCheckInputs = [ pytestCheckHook ];
-  doCheck = torch.cudaSupport;
+  pythonImportsCheck = [ "skrl" ];
+
+  nativeCheckInputs = [
+    flax
+    jax
+    optax
+    pettingzoo
+    pygame
+    pymunk
+    pytestCheckHook
+  ];
+
+  disabledTests = [
+    # TypeError: The array passed to from_dlpack must have __dlpack__ and __dlpack_device__ methods
+    "test_env"
+    "test_multi_agent_env"
 
-  pythonImportsCheck = [
-    "skrl"
-    "skrl.agents"
-    "skrl.agents.torch"
-    "skrl.envs"
-    "skrl.envs.torch"
-    "skrl.models"
-    "skrl.models.torch"
-    "skrl.resources"
-    "skrl.resources.noises"
-    "skrl.resources.noises.torch"
-    "skrl.resources.schedulers"
-    "skrl.resources.schedulers.torch"
-    "skrl.trainers"
-    "skrl.trainers.torch"
-    "skrl.utils"
-    "skrl.utils.model_instantiators"
+    # OverflowError
+    "test_key"
   ];
 
   meta = {
     description = "Reinforcement learning library using PyTorch focusing on readability and simplicity";
-    changelog = "https://github.com/Toni-SM/skrl/releases/tag/${version}";
     homepage = "https://skrl.readthedocs.io";
+    downloadPage = "https://github.com/Toni-SM/skrl";
+    changelog = "https://github.com/Toni-SM/skrl/releases/tag/${finalAttrs.src.tag}";
     license = lib.licenses.mit;
     maintainers = with lib.maintainers; [ bcdarwin ];
   };
-}
+})