From 905d0acefeae7677e48a19f1326d547552ee3c58 Mon Sep 17 00:00:00 2001
From: David Crompton <david.crompton@mail.utoronto.ca>
Date: Fri, 10 May 2024 14:38:47 -0400
Subject: [PATCH] Python: Merge packages between jupyter & system

---
 machines/kcnhub/packages/python.nix    | 57 +++++++++++++++++++++++---
 machines/kcnhub/servers/jupyterhub.nix | 47 +--------------------
 2 files changed, 54 insertions(+), 50 deletions(-)

diff --git a/machines/kcnhub/packages/python.nix b/machines/kcnhub/packages/python.nix
index d18088c..1cd1562 100644
--- a/machines/kcnhub/packages/python.nix
+++ b/machines/kcnhub/packages/python.nix
@@ -1,13 +1,60 @@
 {config, lib, pkgs, ...}: {
-  environment.systemPackages = with pkgs; [
-    (pkgs.python3.withPackages (p: with p; [
+  options = {
+    pythonPackages = lib.mkOption {
+      type = lib.types.anything;
+      default = p: with p; [];
+      example = p: with p; [ numpy scipy ];
+      description = "List of Python Packages installed";
+    };
+  };
+  config = {
+    environment.systemPackages = with pkgs; [
+      ((pkgs.python3.withPackages config.pythonPackages).override (args: { ignoreCollisions = true; }))
+    ];
+    pythonPackages = (p: with p; [
+      # Common Python Libraries
       numpy
       scipy
+      scikit-learn
       pandas
       matplotlib
-      tqdm
+      ipympl
+      matplotlib-inline
       seaborn
+
+      # Neural Data Processing Libraries
+      # Temporarily broken?: https://github.com/NixOS/nixpkgs/issues/259812
+      mne-python
+      nibabel
+      # View niftis!
+      ipyniivue
+      neo
+      
+      # Neurosimulators
       neuronpy
-    ]))
-  ];
+      # nest with MPI is incompatible with Jupyter:
+      # https://www.nest-simulator.org/pynest-api/_modules/nest.html
+      # see pynestkernel comments about their workaround
+      (nest.override {
+        withMpi = false;
+      })
+      brian2
+
+      # Machine learning toolkits:
+      tensorflow
+      tensorboard
+      keras
+      edward
+      transformers
+      tflearn
+      torch
+      torchvision
+      torchsde
+      torchmetrics
+      torchio
+      torchdiffeq
+      botorch
+      lion-pytorch
+    ]);
+  };
 }
diff --git a/machines/kcnhub/servers/jupyterhub.nix b/machines/kcnhub/servers/jupyterhub.nix
index 892d177..0c84eb8 100644
--- a/machines/kcnhub/servers/jupyterhub.nix
+++ b/machines/kcnhub/servers/jupyterhub.nix
@@ -175,7 +175,8 @@
     
     kernels = {
       python3 = let
-        env = (pkgs.python3.withPackages (pythonPackages: with pythonPackages; [
+        env = (pkgs.python3.withPackages (pythonPackages: with pythonPackages;
+          (config.pythonPackages pythonPackages) ++ [
           # Necessary for use as a kernel
           ipykernel
           ipdb
@@ -185,52 +186,8 @@
           mediapy
           # Progress bars etc. for in Jupyter/IPython
           halo
-
-          # Common Python Libraries
-          numpy
-          scipy
-          scikit-learn
-          pandas
           # Tables displayed in Jupyter: like Pandas Dataframes
           ipytablewidgets
-          matplotlib
-          ipympl
-          matplotlib-inline
-          seaborn
-
-          # Neural Data Processing Libraries
-          # Temporarily broken?: https://github.com/NixOS/nixpkgs/issues/259812
-          mne-python
-          nibabel
-          # View niftis!
-          ipyniivue
-          neo
-          
-          # Neurosimulators
-          neuronpy
-          # nest with MPI is incompatible with Jupyter:
-          # https://www.nest-simulator.org/pynest-api/_modules/nest.html
-          # see pynestkernel comments about their workaround
-          (nest.override {
-            withMpi = false;
-          })
-          brian2
-
-          # Machine learning toolkits:
-          tensorflow
-          tensorboard
-          keras
-          edward
-          transformers
-          tflearn
-          torch
-          torchvision
-          torchsde
-          torchmetrics
-          torchio
-          torchdiffeq
-          botorch
-          lion-pytorch
         ])).override (args: { ignoreCollisions = true; });
         # Odd collision between tensorboard of torch & tensorflow
         # need to resolve later