From f9af3388a7fc3f5bfda713c8daef0fbd90db768a Mon Sep 17 00:00:00 2001
From: Manuel Thalmann <m@nuth.ch>
Date: Sat, 7 Dec 2024 17:49:57 +0100
Subject: [PATCH] Store device scripts in `scripts` option

---
 lib/modules/partition/disks.nix | 102 ++++++++++++++++++--------------
 1 file changed, 56 insertions(+), 46 deletions(-)

diff --git a/lib/modules/partition/disks.nix b/lib/modules/partition/disks.nix
index 651a7699..1636d867 100644
--- a/lib/modules/partition/disks.nix
+++ b/lib/modules/partition/disks.nix
@@ -61,9 +61,16 @@ let
           default = { };
         };
 
-        script = mkOption {
-          type = types.str;
-          description = "The script for formatting the disk.";
+        scripts = {
+          init = mkOption {
+            type = types.str;
+            description = "A script for loading the device path into the device variable";
+          };
+
+          partition = mkOption {
+            type = types.str;
+            description = "A script for partitioning and formatting the device.";
+          };
         };
       };
 
@@ -172,35 +179,37 @@ let
           id = "disk-${name}";
           deviceVariable = deviceVar;
 
-          deviceScript =
-            if config.path == null then ''
-              ${deviceSelector}
-            '' else ''
-              ${deviceVarName}=${config.path}
-              if [ ! -b ${deviceVar} ]; then
-                function fallback() {
-                  echo "Couldn't find the specified disk \"${deviceVar}\"."
-                  if fish ${./confirm.fish} "Do you want to choose a different \"${name}\" disk?"; then
-                    ${deviceSelector}
-                  else
-                    exit 1
-                  fi
-                }
+          scripts = {
+            init =
+              if config.path == null then ''
+                ${deviceSelector}
+              '' else ''
+                ${deviceVarName}=${config.path}
+                if [ ! -b ${deviceVar} ]; then
+                  function fallback() {
+                    echo "Couldn't find the specified disk \"${deviceVar}\"."
+                    if fish ${./confirm.fish} "Do you want to choose a different \"${name}\" disk?"; then
+                      ${deviceSelector}
+                    else
+                      exit 1
+                    fi
+                  }
 
-                fallback
-              fi
+                  fallback
+                fi
+              '';
+
+            partition = lib.mkDefault ''
+              function partition() {
+                ${if (!config.wipe) then cleanup else ""}
+                ${probeScript}
+                ${fdiskCommands}
+                ${fixType}
+              }
+
+              partition
             '';
-
-          script = lib.mkDefault ''
-            function partition() {
-              ${if (!config.wipe) then cleanup else ""}
-              ${probeScript}
-              ${fdiskCommands}
-              ${fixType}
-            }
-
-            partition
-          '';
+          };
         };
     }
   );
@@ -332,8 +341,8 @@ in
             partitions = (builtins.concatMap (_: (builtins.attrValues _.partitions)) disks);
           in
           {
-            init = lib.strings.concatLines (builtins.map (_: _.deviceScript) disks);
-            partition = lib.strings.concatLines (builtins.map (_: _.script) disks);
+            init = lib.strings.concatLines (builtins.map (_: _.scripts.init) disks);
+            partition = lib.strings.concatLines (builtins.map (_: _.scripts.partition) disks);
             mount = lib.strings.concatLines (
               (builtins.concatMap
                 (
@@ -341,26 +350,27 @@ in
                     probeScript
                     (builtins.concatStringsSep " " (
                       [ "sudo" "mount" "--mkdir" ] ++
-                        (lib.optionals (_.format == "ntfs") [ "-t" "ntfs3" ]) ++
-                        [
-                          (builtins.concatStringsSep " " (builtins.map (_: "-o ${_}") _.mountOptions))
-                          (partPath _)
-                          (normalizePath "/${rootDir}/${_.mountPoint}")
-                        ]
+                      (lib.optionals (_.format == "ntfs") [ "-t" "ntfs3" ]) ++
+                      [
+                        (builtins.concatStringsSep " " (builtins.map (_: "-o ${_}") _.mountOptions))
+                        (partPath _)
+                        (normalizePath "/${rootDir}/${_.mountPoint}")
+                      ]
                     ))
                   ]
                 )
                 (lib.lists.sortOn
                   (_: normalizePath "/${_.mountPoint}")
                   (builtins.filter (_: _.mountPoint != null) partitions))) ++
-            (builtins.map
-              (
-                _: ''
-                  ${probeScript}
-                  sudo swapon ${partPath _}
-                ''
-              )
-              (builtins.filter (_: _.useSwap) partitions)));
+              (builtins.map
+                (
+                  _: ''
+                    ${probeScript}
+                    sudo swapon ${partPath _}
+                  ''
+                )
+                (builtins.filter (_: _.useSwap) partitions))
+            );
           };
       };
     };