From c78d9e9e24f57cb45100516d2e4ace1c42dce36b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Moritz=20B=C3=B6hme?= <mail@moritzboeh.me>
Date: Thu, 26 Jan 2023 10:55:47 +0100
Subject: [PATCH] rework python module

---
 modules/profiles/desktop.nix |  2 +-
 modules/programs/python.nix  | 66 ++++++++++++++++++++++--------------
 2 files changed, 41 insertions(+), 27 deletions(-)

diff --git a/modules/profiles/desktop.nix b/modules/profiles/desktop.nix
index 687fd40..ffdd4e6 100644
--- a/modules/profiles/desktop.nix
+++ b/modules/profiles/desktop.nix
@@ -78,7 +78,7 @@ with lib; {
       hub.enable = true;
       ledger.enable = true;
       logseq.enable = true;
-      python.enable = true;
+      python.versions."311".enable = true;
       ssh = {
         enable = true;
         includeSecrets = [ ../../secrets/ssh-home.age ];
diff --git a/modules/programs/python.nix b/modules/programs/python.nix
index f0dadfb..adecd55 100644
--- a/modules/programs/python.nix
+++ b/modules/programs/python.nix
@@ -8,41 +8,55 @@ with lib;
 let
   cfg = config.my.programs.python;
 
-  mkPython = packages: version: pkgs.${version}.withPackages (ps: map (flip getAttr ps) packages);
+  pythonVersions = map (version: "3${toString version}") (range 8 11);
+  enabledVersions = filterAttrs (name: value: value.enable) cfg.versions;
 
-  pythonVersions = [
-    "python311"
-    "python310"
-    "python39"
-    "python38"
-  ];
+  pythonPackages = version: attrNames pkgs."python${version}Packages";
 
-  packageLists = map (version: attrNames pkgs."${version}Packages") cfg.versions;
+  commonPackages =
+    let
+      packageLists = map pythonPackages (attrNames enabledVersions);
+    in
+    foldl' intersectLists (head packageLists) (tail packageLists);
 
-  commonPackages = foldl' intersectLists (head packageLists) (tail packageLists);
+  versionOpts = version: {
+    enable = mkEnableOption (toString version);
+    pythonPackages = mkOption {
+      default = [ ];
+      type = with types; listOf (enum (pythonPackages version));
+    };
+  };
 in
 {
   options.my.programs.python = {
-    enable = mkEnableOption "python";
-    versions = mkOption {
-      default = [ "python310" ];
-      type = with types; nonEmptyListOf (enum pythonVersions);
-      example = [ "python39" ];
-    };
-    packages = mkOption {
-      default = [
-        "flake8"
-        "isort"
-        "mypy"
-        "poetry"
-        "pytest"
-        "python-lsp-server"
-      ];
+    versions = genAttrs pythonVersions versionOpts;
+    defaultPackages = mkOption {
+      default = [ ];
       type = with types; listOf (enum commonPackages);
     };
+    extraPackages = mkOption {
+      default = with pkgs; [
+        poetry
+        ruff
+        python310Packages.python-lsp-server
+      ];
+      type = with types; listOf package;
+    };
   };
 
-  config = mkIf cfg.enable {
-    users.users.moritz.packages = map (mkPython cfg.packages) cfg.versions;
+  config = {
+    users.users.moritz.packages =
+      (
+        let
+          mkPython = version:
+            let
+              package = pkgs."python${version}";
+              finalPythonPackages = cfg.versions.${version}.pythonPackages ++ cfg.defaultPackages;
+              getPythonPackages = ps: map (flip getAttr ps) finalPythonPackages;
+            in
+            package.withPackages getPythonPackages;
+        in
+        map mkPython (attrNames enabledVersions)
+      ) ++ cfg.extraPackages;
   };
 }