From 00dc6dd7e73037adb482a817817a29264f6326bf Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Fri, 18 Feb 2022 21:24:20 +0000
Subject: [PATCH] python3Packages.ffcv: init at 0.0.3

---
 .../python-modules/ffcv/default.nix           | 53 +++++++++++++++++++
 pkgs/top-level/python-packages.nix            |  2 +
 2 files changed, 55 insertions(+)
 create mode 100644 pkgs/development/python-modules/ffcv/default.nix

diff --git a/pkgs/development/python-modules/ffcv/default.nix b/pkgs/development/python-modules/ffcv/default.nix
new file mode 100644
index 000000000000..2aee132e294f
--- /dev/null
+++ b/pkgs/development/python-modules/ffcv/default.nix
@@ -0,0 +1,53 @@
+{ buildPythonPackage
+, fetchFromGitHub
+, lib
+, libjpeg
+, numba
+, opencv4
+, pandas
+, pkgconfig
+, pytorch-pfn-extras
+, terminaltables
+, tqdm
+}:
+
+buildPythonPackage rec {
+  pname = "ffcv";
+  version = "0.0.3";
+
+  src = fetchFromGitHub {
+    owner = "libffcv";
+    repo = pname;
+    # See https://github.com/libffcv/ffcv/issues/158.
+    rev = "131d56235eca3f1497bb84eeaec82c3434ef25d8";
+    sha256 = "0f7q2x48lknnf98mqaa35my05qwvdgv0h8l9lpagdw6yhx0a6p2x";
+  };
+
+  # See https://github.com/libffcv/ffcv/issues/159.
+  postPatch = ''
+    substituteInPlace setup.py \
+      --replace "'assertpy'," "" \
+      --replace "'fastargs'," "" \
+      --replace "'imgcat'," "" \
+      --replace "'matplotlib'," "" \
+      --replace "'psutil'," "" \
+      --replace "'sklearn'," "" \
+      --replace "'webdataset'," ""
+  '';
+
+  buildInputs = [ libjpeg pkgconfig ];
+  propagatedBuildInputs = [ opencv4 numba pandas pytorch-pfn-extras terminaltables 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; {
+    description = "FFCV: Fast Forward Computer Vision";
+    homepage = "https://ffcv.io";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ samuela ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index f918f0fe2f1e..f6a82b4530d2 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -2821,6 +2821,8 @@ in {
     hdf5 = pkgs.hdf5_1_10;
   };
 
+  ffcv = callPackage ../development/python-modules/ffcv { };
+
   ffmpeg-python = callPackage ../development/python-modules/ffmpeg-python { };
 
   ffmpeg-progress-yield = callPackage ../development/python-modules/ffmpeg-progress-yield { };