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

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchFromGitHub,
   5
   6   # build-system
   7   flit-core,
   8
   9   # dependencies
  10   numpy,
  11
  12   # optional-dependencies
  13   ipython,
  14   jax,
  15   palettable,
  16
  17   # tests
  18   absl-py,
  19   jaxlib,
  20   pytestCheckHook,
  21   torch,
  22 }:
  23
  24 buildPythonPackage rec {
  25   pname = "treescope";
  26   version = "0.1.6";
  27   pyproject = true;
  28
  29   src = fetchFromGitHub {
  30     owner = "google-deepmind";
  31     repo = "treescope";
  32     rev = "refs/tags/v${version}";
  33     hash = "sha256-QlCKdsQk9VzRNnQKqWUIt6drocx++Aq34cMmZTw0UZw=";
  34   };
  35
  36   build-system = [ flit-core ];
  37
  38   dependencies = [ numpy ];
  39
  40   optional-dependencies = {
  41     notebook = [
  42       ipython
  43       jax
  44       palettable
  45     ];
  46   };
  47
  48   pythonImportsCheck = [ "treescope" ];
  49
  50   nativeCheckInputs = [
  51     absl-py
  52     jax
  53     jaxlib
  54     pytestCheckHook
  55     torch
  56   ];
  57
  58   meta = {
  59     description = "An interactive HTML pretty-printer for machine learning research in IPython notebooks";
  60     homepage = "https://github.com/google-deepmind/treescope";
  61     changelog = "https://github.com/google-deepmind/treescope/releases/tag/v${version}";
  62     license = lib.licenses.asl20;
  63     maintainers = with lib.maintainers; [ GaetanLepage ];
  64   };
  65 }