From 9a8e56a767456d710e036d244633a7feedc394a1 Mon Sep 17 00:00:00 2001
From: jD91mZM2 <me@krake.one>
Date: Fri, 23 Oct 2020 17:54:01 +0200
Subject: [PATCH] Implement support for flakes

---
 flake.nix          | 19 ++++++++++++++-----
 nixos-generate     | 30 ++++++++++++++++++++++++++----
 nixos-generate.nix | 31 ++++++++++++++++++++++---------
 3 files changed, 62 insertions(+), 18 deletions(-)

diff --git a/flake.nix b/flake.nix
index 2696988..69e5f17 100644
--- a/flake.nix
+++ b/flake.nix
@@ -16,25 +16,34 @@
         installFlags = [ "PREFIX=$(out)" ];
         postFixup = ''
           wrapProgram $out/bin/nixos-generate \
-            --prefix PATH : ${pkgs.lib.makeBinPath (with pkgs; [ jq coreutils findutils nix ])}
+            --prefix PATH : ${pkgs.lib.makeBinPath (with pkgs; [ jq coreutils findutils ])}
         '';
       };
+
+      # Currently, you need to mark your configurations with makeOverridable in
+      # order to use nixos-generate on them.
+      nixosConfigurations.example = nixpkgs.lib.makeOverridable nixpkgs.lib.nixosSystem {
+        inherit system;
+        modules = [
+          ./configuration.nix
+        ];
+      };
     });
     defaultPackage = forAllSystems (system: self.packages."${system}".nixos-generators);
 
     devShell = forAllSystems (system: let
       pkgs = nixpkgs.legacyPackages."${system}";
     in pkgs.mkShell {
-      buildInputs = with pkgs; [ jq coreutils findutils nix ];
+      buildInputs = with pkgs; [ jq coreutils findutils ];
     });
 
     # Make it runnable with `nix app`
     apps = forAllSystems (system: {
-      nixos-generators = {
+      nixos-generate = {
         type    = "app";
-        program = "${self.packages."${system}".nixos-generators}/bin/nixos-generators";
+        program = "${self.packages."${system}".nixos-generators}/bin/nixos-generate";
       };
     });
-    defaultApp = forAllSystems (system: self.apps."${system}".nixos-generators);
+    defaultApp = forAllSystems (system: self.apps."${system}".nixos-generate);
   };
 }
diff --git a/nixos-generate b/nixos-generate
index 0a5d0e8..0eaadb6 100755
--- a/nixos-generate
+++ b/nixos-generate
@@ -7,6 +7,8 @@ readonly libexec_dir="${0%/*}"
 readonly format_dir=$libexec_dir/formats
 
 configuration=${NIXOS_CONFIG:-$libexec_dir/configuration.nix}
+flake_uri=
+flake_attr=
 format_path=
 target_system=
 cores=
@@ -28,6 +30,8 @@ Options:
 * --help: shows this help
 * -c, --configuration PATH:
     select the nixos configuration to build. Default: $configuration
+* --flake URI:
+    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
@@ -60,6 +64,16 @@ while [[ $# -gt 0 ]]; do
       configuration=$2
       shift
       ;;
+    --flake)
+      # Note: The reason I'm using awk over cut is because cutting with an
+      # out-of-bounds field will return the last in-bound field instead of empty
+      # string.
+      flake="$(echo "$2" | awk -F'#' '{ print $1; }')"
+
+      flake_uri="$(nix flake info --json -- "$flake" | jq -r .url)"
+      flake_attr="$(echo "$2" | awk -F'#' '{ print $2; }')"
+      shift
+      ;;
     --cores)
       cores=$2
       shift
@@ -122,10 +136,18 @@ if [[ ! -f $format_path ]]; then
   abort "format file not found: $format_path"
 fi
 
-nix_args+=(
-  -I "nixos-config=$configuration"
-  -I "format-config=$format_path"
-)
+nix_args+=(--argstr formatConfig "$(realpath "$format_path")")
+
+if [[ -z "$flake_uri" ]]; then
+  nix_args+=(
+    -I "nixos-config=$configuration"
+  )
+else
+  nix_args+=(
+    --argstr flakeUri "$flake_uri"
+    --argstr flakeAttr "${flake_attr:-"$(hostname)"}"
+  )
+fi
 
 if [[ -n $target_system ]]; then
   nix_args+=(--argstr system "$target_system")
diff --git a/nixos-generate.nix b/nixos-generate.nix
index 109bf81..43856cc 100644
--- a/nixos-generate.nix
+++ b/nixos-generate.nix
@@ -1,7 +1,11 @@
 { nixpkgs ? <nixpkgs>
 , configuration ? <nixos-config>
-, format-config ? <format-config>
 , system ? builtins.currentSystem
+
+, formatConfig
+
+, flakeUri ? null
+, flakeAttr ? null
 }:
 let
   module = { lib, ... }: {
@@ -17,12 +21,21 @@ let
       };
     };
   };
+
+  # Will only get evaluated when used, so no worries
+  flake = builtins.getFlake flakeUri;
+  flakeSystem = flake.outputs.packages."${system}".nixosConfigurations."${flakeAttr}" or flake.outputs.nixosConfigurations."${flakeAttr}";
 in
-import "${toString nixpkgs}/nixos/lib/eval-config.nix" {
-  inherit system;
-  modules = [
-    module
-    format-config
-    configuration
-  ];
-}
+  if flakeUri != null then
+    flakeSystem.override (attrs: {
+      modules = attrs.modules ++ [ module formatConfig ];
+    })
+  else
+    import "${toString nixpkgs}/nixos/lib/eval-config.nix" {
+      inherit system;
+      modules = [
+        module
+        formatConfig
+        configuration
+      ];
+    }