Split up disk script into multiple scripts

lib/modules/partition/disks.nix

@@ -298,64 +298,37 @@ in
             default = { };
           };
 
+          scripts = {
+            init = mkOption {
+              type = types.str;
+              description = "The script for initializing the disk partitioning script.";
+            };
+
+            partition = mkOption {
+              type = types.str;
+              description = "The script for partitioning the disks.";
+            };
+
+            mount = mkOption {
+              type = types.str;
+              description = "The script for mounting the partitioned disks.";
+            };
+          };
+
           script = mkOption {
             type = types.str;
             description = "The script for partitioning the system's disks.";
-          };
+            default =
-        };
-      };
-    };
-  };
-
-  config = {
-    valhalla = {
-      fileSystems.diskSetup = {
-        script = lib.mkDefault (
               let
-            cfg = config.valhalla.fileSystems;
+                cfg = config.valhalla.fileSystems.diskSetup;
-            inherit (cfg) rootDir;
+                inherit (cfg) scripts;
-            inherit (lib.strings) normalizePath;
+                disks = (builtins.attrValues cfg.disks);
-            partPath = part: "/dev/disk/by-label/${part.label}";
-            disks = ((builtins.attrValues cfg.diskSetup.disks));
-            partitions = (builtins.concatMap (_: (builtins.attrValues _.partitions)) disks);
-
-            mountScript = lib.strings.concatLines (builtins.concatMap
-              (
-                _: [
-                  probeScript
-                  (builtins.concatStringsSep " " ([
-                    "sudo"
-                    "mount"
-                    "--mkdir"
-                  ] ++ (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)));
-
-            swapScript = lib.strings.concatLines (builtins.map
-              (
-                _: ''
-                  ${probeScript}
-                  sudo swapon ${partPath _}
-                ''
-              )
-              (builtins.filter (_: _.useSwap) partitions));
               in
-          lib.strings.concatLines ([
+              ''
-            "#!/bin/bash"
+                #!/bin/bash
-            "set -o errexit"
+                set -o errexit
-          ]
+                ${scripts.init}
-          ++ (builtins.map (_: _.deviceScript) disks)
+                ${lib.strings.concatLines (lib.optionals ((builtins.length disks) > 0) [
-          ++ lib.optionals ((builtins.length disks) > 0) [
                   ''echo "$(tput setaf 3)==== WARNING ====$(tput sgr0)"''
                   (''echo "Continuing this script will alter the partitions of ''
                     + (lib.strings.concatStringsSep ", " (builtins.map (_: "${_.deviceVariable}") (lib.lists.init disks)))
@@ -365,11 +338,59 @@ in
                       exit 1
                     fi
                   ''
-          ] ++ (builtins.map (_: _.script) disks) ++ [
+                ])}
-            mountScript
+                ${scripts.partition}
-            swapScript
+                ${scripts.mount}
-          ])
+              '';
-        );
+          };
+        };
+      };
+    };
+  };
+
+  config = {
+    valhalla = {
+      fileSystems.diskSetup = {
+        scripts =
+          let
+            cfg = config.valhalla.fileSystems;
+            inherit (cfg) rootDir;
+            inherit (lib.strings) normalizePath;
+            partPath = part: "/dev/disk/by-label/${part.label}";
+            disks = ((builtins.attrValues cfg.diskSetup.disks));
+            partitions = (builtins.concatMap (_: (builtins.attrValues _.partitions)) disks);
+          in
+          {
+            init = lib.strings.concatLines (builtins.map (_: _.deviceScript) disks);
+            partition = lib.strings.concatLines (builtins.map (_: _.script) disks);
+            mount = lib.strings.concatLines (
+              (builtins.concatMap
+                (
+                  _: [
+                    probeScript
+                    (builtins.concatStringsSep " " (
+                      [ "sudo" "mount" "--mkdir" ] ++
+                        (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)));
+          };
       };
     };
   };