Support additional format search paths

README.md

@@ -207,7 +207,7 @@ images can be built from that flake by running:
 `nixos-generators` can be included as a `Flake` input and provides
 a `nixos-generate` function for building images as `Flake` outputs. This
 approach pins all dependencies and allows for conveniently defining multiple
-output types based on one config. 
+output types based on one config.
 
 An example `flake.nix` demonstrating this approach is below. `vmware` or
 `virtualbox` images can be built from the same `configuration.nix` by running
@@ -237,14 +237,14 @@ multiple custom formats.  `nixosGenerate` will then match against these custom f
           # ./configuration.nix
         ];
         format = "vmware";
-        
+
         # optional arguments:
         # explicit nixpkgs and lib:
         # pkgs = nixpkgs.legacyPackages.x86_64-linux;
         # lib = nixpkgs.legacyPackages.x86_64-linux.lib;
         # additional arguements to pass to modules:
         # specialArgs = { myExtraArg = "foobar"; };
-        
+
         # you can also define your own custom formats
         # customFormats = { "myFormat" = <myFormatModule>; ... };
         # format = "myFormat";
@@ -268,6 +268,46 @@ multiple custom formats.  `nixosGenerate` will then match against these custom f
 * If boot fails for some reason, you will not get a recovery shell unless the root user is enabled, which you can do by setting a password for them (`users.users.root.password = "something";`, possibly `users.mutableUsers = true;` so you can interactively change the passwords after boot)
 * After booting, if you intend to use `nixos-switch`, consider using `nixos-generate-config`.
 
+## Using custom formats
+
+You can choose a format by telling `nixos-generate` its full path:
+
+```console
+nixos-generate --format-path ./path/to/my-format.nix
+```
+
+Additionally, you can tell `nixos-generate` where to search for format files by
+
+* Adding `:`-separated paths to the `NIXOS_GENERATORS_FORMAT_SEARCH_PATH`
+  environment variable, or
+* Calling `nixos-generate` with one or more `--format-search-path <path>`
+  options.
+
+Example:
+
+```console
+NIXOS_GENERATORS_FORMAT_SEARCH_PATH=/path/a:/path/b nixos-generate --format-search-path /path/c --format-search-path /path/d -f my-format
+```
+
+The above command searches for the file `my-format.nix` in the following paths,
+in order from highest precedence to lowest:
+
+1. `/path/d/my-format.nix`
+2. `/path/c/my-format.nix`
+3. `/path/a/my-format.nix`
+4. `/path/b/my-format.nix`
+5. `my-format.nix` in the builtin `nixos-generate` format directory
+
+Note that:
+
+* `nixos-generate` does not recognize a mechanism for escaping `:` characters
+  in paths specified in `NIXOS_GENERATORS_FORMAT_SEARCH_PATH`; if you have
+  custom formats that live in a path that contains `:`, specify the path with
+  `--format-search-path ./path/that/contains/a:or/two:`.
+* `nixos-generate` ignores empty strings in the list of format search paths
+  (`nixos-generate --format-search-path ''`).
+* Format names cannot be empty and cannot contain `/` elements.
+
 ### License
 
 This project is licensed under the [MIT License](LICENSE).

checks/fixtures/formats/a/foo.nix

@@ -0,0 +1 @@
+{}

checks/fixtures/formats/b/bar.nix

@@ -0,0 +1 @@
+{}

checks/fixtures/formats/c/baz.nix

@@ -0,0 +1 @@
+{}

checks/fixtures/formats/d/quux.nix

@@ -0,0 +1 @@
+{}

checks/test-format-search-path.nix

@@ -0,0 +1,78 @@
+{
+  nixpkgs,
+  self,
+  system,
+}: let
+  inherit
+    (self.packages.${system})
+    nixos-generate
+    ;
+
+  pkgs = nixpkgs.legacyPackages.${system};
+in
+  pkgs.runCommand "test-format-search-path" {} ''
+    rc=$?
+
+    fail() {
+      rc="$?"
+
+      if (( "$#" > 0 )); then
+        printf 1>&2 -- "$@"
+      else
+        printf 1>&2 -- 'unknown error\n'
+      fi
+    }
+
+    run() {
+      ${nixos-generate}/bin/nixos-generate "$@"
+    }
+
+    showFormatSearchPath() {
+      run --show-format-search-path "$@"
+    }
+
+    list() {
+      run --list "$@"
+    }
+
+    fixtures="${self}/checks/fixtures/formats"
+    builtin="${nixos-generate}/share/nixos-generator/formats"
+
+    path="$(showFormatSearchPath)" || fail 'error running nixos-generate\n'
+
+    expected="$builtin"
+
+    [[ "$path" == "$expected" ]] \
+      || fail 'expected format search path to contain:\n%s\ngot:\n%s\n' "$expected" "$path"
+
+    export NIXOS_GENERATORS_FORMAT_SEARCH_PATH="''${fixtures}/c:''${fixtures}/d"
+
+    path="$(showFormatSearchPath --format-search-path "''${fixtures}/b" --format-search-path "''${fixtures}/a")" \
+      || fail 'error running nixos-generate\n'
+
+    expected="\
+    ''${fixtures}/a
+    ''${fixtures}/b
+    ''${fixtures}/c
+    ''${fixtures}/d
+    $builtin"
+
+    [[ "$path" == "$expected" ]] \
+      || fail 'expected format search path to contain:\n%s\ngot:\n%s\n' "$expected" "$path"
+
+    declare -A formats
+    while read -r format; do
+      formats["$format"]=1
+    done < <(list --format-search-path "''${fixtures}/b" --format-search-path "''${fixtures}/a")
+
+    for format in foo bar baz quux; do
+      [[ -n "''${formats["$format"]:-}" ]] \
+        || fail 'expected formats to include %s\n' "$format"
+    done
+
+    if (( rc == 0 )); then
+      touch "$out"
+    fi
+
+    exit "$rc"
+  ''

flake.nix

@@ -151,6 +151,9 @@
                 test-customize-format = import ./checks/test-customize-format.nix {
                   inherit nixpkgs system;
                 };
+                test-format-search-path = import ./checks/test-format-search-path.nix {
+                  inherit nixpkgs self system;
+                };
               in
                 lib.mapAttrs makeLazyDrv (
                   {
@@ -159,7 +162,7 @@
                       nixos-generate
                       ;
 
-                    inherit test-customize-format;
+                    inherit test-customize-format test-format-search-path;
 
                     is-formatted = import ./checks/is-formatted.nix {
                       pkgs = nixpkgs.legacyPackages.${system};

nixos-generate

@@ -4,21 +4,27 @@ set -euo pipefail
 ## Configuration
 
 readonly libexec_dir="${0%/*}"
-readonly format_dir=$libexec_dir/formats
 
 configuration=${NIXOS_CONFIG:-$libexec_dir/configuration.nix}
 flake_uri=
 flake_attr=
+format=
 format_path=
 target_system=
 cores=
 run=
+list_formats=false
+show_format_search_path=false
 nix_args=(
   "$libexec_dir/nixos-generate.nix"
 )
 has_outlink=false
 nix_build_args=()
 
+# `printf' rather than `<<<' to avoid introducing a spurious trailing newline
+mapfile -t -d : format_dirs < <(printf -- '%s' "${NIXOS_GENERATORS_FORMAT_SEARCH_PATH:-}")
+format_dirs+=("$libexec_dir/formats")
+
 ## Functions
 
 showUsage() {
@@ -34,7 +40,11 @@ Options:
     selects the nixos configuration to build, using flake uri like "~/dotfiles#my-config"
 * -f, --format NAME: select one of the pre-determined formats
 * --format-path PATH: pass a custom format
-* --list: list the available built-in formats
+* --format-search-path DIR:
+    prepend a directory to the list of directories ${0##*/} searches for format definitions
+* --list: list the available formats
+* --show-format-search-path:
+    list the directories ${0##*/} searches for format files
 * --run: runs the configuration in a VM
          only works for the "vm" and "vm-nogui" formats
 * --show-trace: show more detailed nix evaluation location information
@@ -47,9 +57,67 @@ USAGE
 }
 
 listFormats() {
-  for format in "$format_dir"/*.nix; do
-    basename "$format" ".nix"
+  local -A formats
+  local format_dir format_file format
+
+  for format_dir in "${format_dirs[@]}"; do
+    if [[ -n $format_dir ]]; then
+      for format_file in "$format_dir"/*.nix; do
+          if [[ -f "$format_file" ]]; then
+            format=$(basename "$format_file" ".nix")
+            formats["$format"]=1
+          fi
+      done
+    fi
+  done
+
+  for format in "${!formats[@]}"; do
+    printf -- '%s\n' "$format"
+  done | sort
+}
+
+showFormatSearchPath() {
+  local format_dir
+
+  for format_dir in "${format_dirs[@]}"; do
+    if [[ -n $format_dir ]]; then
+      printf -- '%s\n' "$format_dir"
+    fi
+  done
+}
+
+validateFormat() {
+  case "${1:-}" in
+    */* | '')
+      abort "not a valid format name: ${1:-<empty>}"
+      return 1
+      ;;
+  esac
+}
+
+findFormat() {
+  local format="${1?}"
+  shift
+
+  validateFormat "$format" || return
+
+  local -n ref_var="${1:-format_file}"
+  shift
+
+  local format_dir maybe_format_file
+
+  for format_dir in "${format_dirs[@]}"; do
+    if [[ -n $format_dir ]]; then
+      maybe_format_file="${format_dir}/${format}.nix"
+
+      if [[ -f "$maybe_format_file" ]]; then
+        ref_var="$maybe_format_file"
+        return
+      fi
+    fi
   done
+
+  abort "unable to locate file for format: $format"
 }
 
 abort() {
@@ -84,27 +152,33 @@ while [[ $# -gt 0 ]]; do
       shift 2
       ;;
     -f | --format)
-      format_path=$format_dir/$2.nix
+      format="$2"
       shift
       ;;
     --format-path)
       format_path=$2
       shift
       ;;
+    --format-search-path)
+      format_dirs=("$2" "${format_dirs[@]}")
+      shift
+      ;;
     --help)
       showUsage
       exit
       ;;
     --list)
-      listFormats
-      exit
+      list_formats=true
+      show_format_search_path=false
+      ;;
+    --show-format-search-path)
+      list_formats=false
+      show_format_search_path=true
       ;;
     --run)
       run=1
       # default to the VM format
-      if [[ -z $format_path ]]; then
-        format_path=$format_dir/vm.nix
-      fi
+      format="${format:-vm}"
       ;;
     --show-trace)
       nix_args+=(--show-trace)
@@ -129,12 +203,24 @@ while [[ $# -gt 0 ]]; do
   shift
 done
 
+if $list_formats; then
+  listFormats
+  exit
+elif $show_format_search_path; then
+  showFormatSearchPath
+  exit
+fi
+
 if ! $has_outlink; then
   nix_build_args+=(--no-out-link)
 fi
 
 if [[ -z $format_path ]]; then
-  abort "missing format. use --help for more details"
+  if [[ -n $format ]] ;then
+    findFormat "$format" format_path
+  else
+    abort "missing format. use --help for more details"
+  fi
 fi
 
 if [[ ! -f $format_path ]]; then