python subpackage options for extras

Author: thomasjm

modules/kernels/python/default.nix

@@ -9,6 +9,7 @@
 
 , settings
 , settingsSchema
+, subPackageSettingsSchema
 }:
 
 with { inherit (settings) packages; };
@@ -23,11 +24,31 @@ let
   displayName = name + " " +  python3.version;
   kernelName = "python3";
 
-  packageOptions = python3.pkgs;
-  packageSearch = common.searcher packageOptions;
+  packageOptionsBase = python3.pkgs;
+  packageSearch = common.searcher packageOptionsBase;
+
+  # For makeEnvironment.nix's chooseMeta lookup: lazily add settingsSchema to packages
+  # that have optional-dependencies, so the UI knows extras are available.
+  # Nix's laziness means only actually-accessed packages get evaluated.
+  packageOptions = lib.mapAttrs (_name: pkg:
+    if (builtins.tryEval (pkg ? optional-dependencies && pkg.optional-dependencies != {})).value or false
+    then pkg // { inherit subPackageSettingsSchema; settingsSchema = subPackageSettingsSchema; }
+    else pkg
+  ) packageOptionsBase;
+
+  resolvePackage = p:
+    let
+      pname = common.packageName p;
+      basePkg = builtins.getAttr pname packageOptionsBase;
+      extras = if lib.isAttrs p then (p.extras or []) else [];
+      extraPkgs = lib.concatMap (extra:
+        basePkg.optional-dependencies.${extra} or []
+      ) extras;
+    in [basePkg] ++ extraPkgs;
+
   allPackages =
-    [packageOptions.ipykernel packageOptions.ipywidgets]
-    ++ map (x: builtins.getAttr x packageOptions) (map common.packageName packages)
+    [packageOptionsBase.ipykernel packageOptionsBase.ipywidgets]
+    ++ lib.concatMap resolvePackage packages
   ;
 
   pythonWithPackages = f: python3.withPackages (_: allPackages ++ f packageOptions);

modules/kernels/python/module.nix

@@ -13,9 +13,30 @@ let
         example = "Package outputs to include";
         type = types.listOf types.str;
       };
+      extras = mkOption {
+        example = "Extras (optional-dependencies) to enable for the package";
+        type = types.listOf types.str;
+        default = [];
+      };
     };
   };
 
+  subPackageEvaluated = lib.evalModules {
+    modules = [
+      {
+        options = {
+          subPackage = {
+            extras = mkOption {
+              example = "Extras (optional-dependencies) to enable for the package";
+              type = types.listOf types.str;
+              default = [];
+            };
+          };
+        };
+      }
+    ];
+  };
+
   mkOptions = packageOption: {
     enable = mkOption {
       example = "Enable Python 3 kernel";
@@ -162,6 +183,7 @@ in
 
           settings = config.kernels.python3;
           settingsSchema = nixosOptionsToSettingsSchema { componentsToDrop = 2; } options.kernels.python3;
+          subPackageSettingsSchema = nixosOptionsToSettingsSchema { componentsToDrop = 1; } subPackageEvaluated.options.subPackage;
         };
     })
 
@@ -174,6 +196,7 @@ in
 
           settings = config.kernels.pypy3;
           settingsSchema = nixosOptionsToSettingsSchema { componentsToDrop = 2; } options.kernels.pypy3;
+          subPackageSettingsSchema = nixosOptionsToSettingsSchema { componentsToDrop = 1; } subPackageEvaluated.options.subPackage;
         };
     })
   ];