python3Packages.ffcv: fix tests

Fixes python3Packages.ffcv tests by changing the directory in precheck.
Cleaned up and refactored the derivation.

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

@@ -1,6 +1,8 @@
 {
   buildPythonPackage,
+  pythonAtLeast,
   fetchFromGitHub,
+  setuptools,
   lib,
   libjpeg,
   numba,
@@ -10,12 +12,20 @@
   pytorch-pfn-extras,
   terminaltables,
   tqdm,
+  pytestCheckHook,
+  assertpy,
+  psutil,
+  torchvision,
+  webdataset,
 }:
 
 buildPythonPackage rec {
   pname = "ffcv";
   version = "1.0.0";
-  format = "setuptools";
+  pyproject = true;
+
+  # version 1.0.0 uses distutils which was removed in Python 3.12
+  disabled = pythonAtLeast "3.12";
 
   src = fetchFromGitHub {
     owner = "libffcv";
@@ -27,15 +37,13 @@ buildPythonPackage rec {
   # See https://github.com/libffcv/ffcv/issues/159.
   postPatch = ''
     substituteInPlace setup.py \
-      --replace "'assertpy'," "" \
+      --replace-fail "'assertpy'," "" \
-      --replace "'fastargs'," "" \
+      --replace-fail "'fastargs'," "" \
-      --replace "'imgcat'," "" \
+      --replace-fail "'opencv-python'," "" \
-      --replace "'matplotlib'," "" \
+      --replace-fail "'psutil'," "" \
-      --replace "'psutil'," "" \
-      --replace "'sklearn'," "" \
-      --replace "'webdataset'," ""
   '';
 
+  build-system = [ setuptools ];
   nativeBuildInputs = [ pkg-config ];
   buildInputs = [ libjpeg ];
   propagatedBuildInputs = [
@@ -47,16 +55,58 @@ buildPythonPackage rec {
     tqdm
   ];
 
-  # `ffcv._libffcv*.so` cannot be loaded in the nix build environment for some
-  # reason. See https://github.com/NixOS/nixpkgs/pull/160441#issuecomment-1045204722.
-  doCheck = false;
-
   pythonImportsCheck = [ "ffcv" ];
 
-  meta = with lib; {
+  # C/C++ python modules are only in the installed output and not in the build
+  # directory. Since tests are run from the build directory python prefers to
+  # import the local module first which does not contain the C/C++ python
+  # modules and results in an import error. By changing the directory to
+  # 'tests' the build directory is no long available and python will import
+  # from the installed output in the nix store which does contain the C/C++
+  # python modules.
+  preCheck = ''
+    cd tests
+  '';
+
+  nativeCheckInputs = [
+    assertpy
+    psutil
+    pytestCheckHook
+    torchvision
+    webdataset
+  ];
+
+  disabledTestPaths = [
+    # Tests require network access and do not work in the sandbox
+    "test_augmentations.py"
+    # Occasionally causes the testing phase to hang
+    "test_basic_pipeline.py"
+  ];
+
+  disabledTests = [
+    # Tests require network access and do not work in the sandbox
+    "test_cifar_subset"
+    # Requires CUDA which is unfree and unfree packages are not built by Hydra
+    "test_cuda"
+    "test_gpu_normalization"
+    # torch.multiprocessing.spawn.ProcessRaisedException
+    "test_traversal_sequential_2"
+    "test_traversal_sequential_3"
+    "test_traversal_sequential_4"
+    "test_traversal_random_2"
+    "test_traversal_random_3"
+    "test_traversal_random_4"
+    "test_traversal_sequential_distributed_with_indices"
+    "test_traversal_random_distributed_with_indices"
+  ];
+
+  meta = {
     description = "FFCV: Fast Forward Computer Vision";
     homepage = "https://ffcv.io";
-    license = licenses.asl20;
+    license = lib.licenses.asl20;
-    maintainers = with maintainers; [ samuela ];
+    maintainers = with lib.maintainers; [
+      samuela
+      djacu
+    ];
   };
 }