From 28d4e7195083c4d3ad370094c17d7cd028156e73 Mon Sep 17 00:00:00 2001
From: lassulus <git@lassul.us>
Date: Sat, 20 May 2023 13:55:25 +0200
Subject: [PATCH] cli: consume diskoConfigurations or nixosConfigurations disko
 configs

---
 cli.nix       | 66 +++++++++++++++++++++++++++++++++------------------
 default.nix   |  6 ++---
 module.nix    | 10 +++++++-
 tests/lib.nix |  2 +-
 4 files changed, 56 insertions(+), 28 deletions(-)

diff --git a/cli.nix b/cli.nix
index 3dedbe1..6af44e6 100644
--- a/cli.nix
+++ b/cli.nix
@@ -14,31 +14,51 @@ let
     inherit lib;
   };
 
-  diskFormat =
-    if flake != null then
-      (lib.attrByPath [ "diskoConfigurations" flakeAttr ] (builtins.abort "${flakeAttr} does not exist") (builtins.getFlake flake)) args
+  diskoAttr =
+    if noDeps then
+      {
+        create = "createScriptNoDeps";
+        mount = "mountScriptNoDeps";
+        zap_create_mount = "diskoNoDeps";
+        disko = "diskoNoDeps";
+      }.${mode}
     else
-      import diskoFile ({ inherit lib; } // args);
+      {
+        create = "createScript";
+        mount = "mountScript";
+        zap_create_mount = "diskoScript";
+        disko = "diskoScript";
+      }.${mode};
+
+  hasDiskoConfigFlake =
+    diskoFile != null || lib.hasAttrByPath [ "diskoConfigurations" flakeAttr ] (builtins.getFlake flake);
+
+  hasDiskoModuleFlake =
+    lib.hasAttrByPath [ "nixosConfigurations" flakeAttr "config" "disko" "devices" ] (builtins.getFlake flake);
+
+  diskFormat =
+    let
+      diskoConfig =
+        if diskoFile != null then
+          import diskoFile
+        else
+          (builtins.getFlake flake).diskoConfigurations.${flakeAttr};
+    in
+    if builtins.isFunction diskoConfig then
+      diskoConfig ({ inherit lib; } // args)
+    else
+      diskoConfig;
 
   diskoEval =
-    if noDeps then
-      if (mode == "create") then
-        disko.createScriptNoDeps diskFormat pkgs
-      else if (mode == "mount") then
-        disko.mountScriptNoDeps diskFormat pkgs
-      else if (mode == "zap_create_mount") then
-        disko.zapCreateMountScriptNoDeps diskFormat pkgs
-      else
-        builtins.abort "invalid mode"
+    disko.${diskoAttr} diskFormat pkgs;
+
+  diskoScript =
+    if hasDiskoConfigFlake then
+      diskoEval
+    else if (lib.traceValSeq hasDiskoModuleFlake) then
+      (builtins.getFlake flake).nixosConfigurations.${flakeAttr}.config.system.build.${diskoAttr}
     else
-      if (mode == "create") then
-        disko.createScript diskFormat pkgs
-      else if (mode == "mount") then
-        disko.mountScript diskFormat pkgs
-      else if (mode == "zap_create_mount") then
-        disko.zapCreateMountScript diskFormat pkgs
-      else
-        builtins.abort "invalid mode"
-  ;
+      (builtins.abort "neither diskoConfigurations.${flakeAttr} nor nixosConfigurations.${flakeAttr} found");
+
 in
-diskoEval
+diskoScript
diff --git a/default.nix b/default.nix
index f70e9d3..dd42ada 100644
--- a/default.nix
+++ b/default.nix
@@ -34,12 +34,12 @@ in
   mountScriptNoDeps = cfg: pkgs: (diskoLib.writeCheckedBash { inherit pkgs checked; noDeps = true; }) "disko-mount" ''
     ${diskoLib.mount (eval cfg).config.disko.devices}
   '';
-  zapCreateMount = cfg: diskoLib.zapCreateMount (eval cfg).config.disko.devices;
-  zapCreateMountScript = cfg: pkgs: (diskoLib.writeCheckedBash { inherit pkgs checked; }) "disko-zap-create-mount" ''
+  disko = cfg: diskoLib.zapCreateMount (eval cfg).config.disko.devices;
+  diskoScript = cfg: pkgs: (diskoLib.writeCheckedBash { inherit pkgs checked; }) "disko-zap-create-mount" ''
     export PATH=${lib.makeBinPath (diskoLib.packages (eval cfg).config.disko.devices pkgs)}:$PATH
     ${diskoLib.zapCreateMount (eval cfg).config.disko.devices}
   '';
-  zapCreateMountScriptNoDeps = cfg: pkgs: (diskoLib.writeCheckedBash { inherit pkgs checked; noDeps = true; }) "disko-zap-create-mount" ''
+  diskoNoDeps = cfg: pkgs: (diskoLib.writeCheckedBash { inherit pkgs checked; noDeps = true; }) "disko-zap-create-mount" ''
     ${diskoLib.zapCreateMount (eval cfg).config.disko.devices}
   '';
   config = cfg: { imports = diskoLib.config (eval cfg).config.disko.devices; };
diff --git a/module.nix b/module.nix
index 1262841..4529b0d 100644
--- a/module.nix
+++ b/module.nix
@@ -52,7 +52,15 @@ in
       ${diskoLib.zapCreateMount cfg.devices}
     '';
 
-    # This is useful to skip copying executables uploading a script to an in-memory installer
+    # These are useful to skip copying executables uploading a script to an in-memory installer
+    system.build.formatScriptNoDeps = (diskoLib.writeCheckedBash { inherit pkgs checked; noDeps = true; }) "disko-create" ''
+      ${diskoLib.create cfg.devices}
+    '';
+
+    system.build.mountScriptNoDeps = (diskoLib.writeCheckedBash { inherit pkgs checked; noDeps = true; }) "disko-mount" ''
+      ${diskoLib.mount cfg.devices}
+    '';
+
     system.build.diskoNoDeps = (diskoLib.writeCheckedBash { inherit pkgs checked; noDeps = true; }) "disko" ''
       ${diskoLib.zapCreateMount cfg.devices}
     '';
diff --git a/tests/lib.nix b/tests/lib.nix
index 2c2a090..50c41d0 100644
--- a/tests/lib.nix
+++ b/tests/lib.nix
@@ -28,7 +28,7 @@
       tsp-generator = pkgs.callPackage ../. { checked = true; };
       tsp-create = (tsp-generator.createScript (import disko-config { disks = builtins.tail disks; inherit lib; })) pkgs;
       tsp-mount = (tsp-generator.mountScript (import disko-config { disks = builtins.tail disks; inherit lib; })) pkgs;
-      tsp-disko = (tsp-generator.zapCreateMountScript (import disko-config { disks = builtins.tail disks; inherit lib; })) pkgs;
+      tsp-disko = (tsp-generator.diskoScript (import disko-config { disks = builtins.tail disks; inherit lib; })) pkgs;
       tsp-config = tsp-generator.config (import disko-config { inherit disks; inherit lib; });
       num-disks = builtins.length (lib.attrNames (import disko-config { inherit lib; }).disko.devices.disk);
       installed-system = { modulesPath, ... }: {