diff --git a/pkgs/development/python-modules/treescope/default.nix b/pkgs/development/python-modules/treescope/default.nix new file mode 100644 index 000000000000..c27e95cdff38 --- /dev/null +++ b/pkgs/development/python-modules/treescope/default.nix @@ -0,0 +1,65 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + flit-core, + + # dependencies + numpy, + + # optional-dependencies + ipython, + jax, + palettable, + + # tests + absl-py, + jaxlib, + pytestCheckHook, + torch, +}: + +buildPythonPackage rec { + pname = "treescope"; + version = "0.1.5"; + pyproject = true; + + src = fetchFromGitHub { + owner = "google-deepmind"; + repo = "treescope"; + rev = "refs/tags/v${version}"; + hash = "sha256-+Hm60O9tEXIiE0av1O0BsOdMln4e1s7ijb3WNiQ74jE="; + }; + + build-system = [ flit-core ]; + + dependencies = [ numpy ]; + + optional-dependencies = { + notebook = [ + ipython + jax + palettable + ]; + }; + + pythonImportsCheck = [ "treescope" ]; + + nativeCheckInputs = [ + absl-py + jax + jaxlib + pytestCheckHook + torch + ]; + + meta = { + description = "An interactive HTML pretty-printer for machine learning research in IPython notebooks"; + homepage = "https://github.com/google-deepmind/treescope"; + changelog = "https://github.com/google-deepmind/treescope/releases/tag/v${version}"; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ GaetanLepage ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 2c4b0027a301..e743b3dc785a 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -15773,6 +15773,8 @@ self: super: with self; { treeo = callPackage ../development/python-modules/treeo { }; + treescope = callPackage ../development/python-modules/treescope { }; + treex = callPackage ../development/python-modules/treex { }; treq = callPackage ../development/python-modules/treq { };