From 8f3bfc39aa668f42f128c6623438dff72ec7eb43 Mon Sep 17 00:00:00 2001
From: Archit Gupta <archit@accelbread.com>
Date: Tue, 5 Dec 2023 19:44:12 -0800
Subject: [PATCH] Allow devShell to be set to a package definition

---
 API_GUIDE.md                 |  5 ++-
 builtinModules/devShells.nix | 81 ++++++++++++++++++++----------------
 tests/default.nix            |  8 +++-
 3 files changed, 56 insertions(+), 38 deletions(-)

diff --git a/API_GUIDE.md b/API_GUIDE.md
index 59bd987..ed66846 100644
--- a/API_GUIDE.md
+++ b/API_GUIDE.md
@@ -371,6 +371,9 @@ such an attribute set in order to access packages.
 `devShell.stdenv` allows changing the stdenv used for the shell. It is a
 function that takes the package set and returns the stdenv to use.
 
+`devShell` can alternatively be set to a package definition, which is then used
+as the default shell, overriding the above options.
+
 For example, these can be configured as follows:
 
 ```nix
@@ -397,7 +400,7 @@ For example, these can be configured as follows:
 
 The above exports `devShells.${system}.default` outputs.
 
-To add build inputs of one of your packages, you can do as follows:
+To add the build inputs of one of your packages, you can do as follows:
 
 ```nix
 {
diff --git a/builtinModules/devShells.nix b/builtinModules/devShells.nix
index 1f0f02e..678df81 100644
--- a/builtinModules/devShells.nix
+++ b/builtinModules/devShells.nix
@@ -5,43 +5,51 @@
 { config, lib, flakelight, ... }:
 let
   inherit (lib) filterAttrs mapAttrs mkDefault mkIf mkMerge mkOption;
-  inherit (lib.types) functionTo lazyAttrsOf lines listOf nullOr package str
-    submodule;
+  inherit (lib.types) coercedTo functionTo lazyAttrsOf lines listOf nullOr
+    package str submodule;
   inherit (flakelight) supportedSystem;
   inherit (flakelight.types) optFunctionTo packageDef;
+
+  devShellModule.options = {
+    inputsFrom = mkOption {
+      type = functionTo (listOf package);
+      default = _: [ ];
+    };
+
+    packages = mkOption {
+      type = functionTo (listOf package);
+      default = _: [ ];
+    };
+
+    shellHook = mkOption {
+      type = optFunctionTo lines;
+      default = "";
+    };
+
+    env = mkOption {
+      type = optFunctionTo (lazyAttrsOf str);
+      default = { };
+    };
+
+    stdenv = mkOption {
+      type = functionTo package;
+      default = pkgs: pkgs.stdenv;
+    };
+
+    overrideShell = mkOption {
+      type = nullOr packageDef;
+      internal = true;
+      default = null;
+    };
+  };
 in
 {
   options = {
     devShell = mkOption {
       default = null;
-      type = nullOr (submodule {
-        options = {
-          inputsFrom = mkOption {
-            type = functionTo (listOf package);
-            default = _: [ ];
-          };
-
-          packages = mkOption {
-            type = functionTo (listOf package);
-            default = _: [ ];
-          };
-
-          shellHook = mkOption {
-            type = optFunctionTo lines;
-            default = "";
-          };
-
-          env = mkOption {
-            type = optFunctionTo (lazyAttrsOf str);
-            default = { };
-          };
-
-          stdenv = mkOption {
-            type = functionTo package;
-            default = pkgs: pkgs.stdenv;
-          };
-        };
-      });
+      type = nullOr (coercedTo packageDef
+        (x: { overrideShell = x; })
+        (submodule devShellModule));
     };
 
     devShells = mkOption {
@@ -53,12 +61,13 @@ in
   config = mkMerge [
     (mkIf (config.devShell != null) {
       devShells.default = mkDefault ({ pkgs, mkShell }:
-        mkShell.override { stdenv = config.devShell.stdenv pkgs; }
-          ((config.devShell.env pkgs) // {
-            inputsFrom = config.devShell.inputsFrom pkgs;
-            packages = config.devShell.packages pkgs;
-            shellHook = config.devShell.shellHook pkgs;
-          }));
+        let cfg = mapAttrs (_: v: v pkgs) config.devShell; in
+        mkShell.override { inherit (cfg) stdenv; }
+          (cfg.env // { inherit (cfg) inputsFrom packages shellHook; }));
+    })
+
+    (mkIf (config.devShell.overrideShell or null != null) {
+      devShells.default = config.devShell.overrideShell;
     })
 
     (mkIf (config.devShells != { }) {
diff --git a/tests/default.nix b/tests/default.nix
index a35158d..2ea4e42 100644
--- a/tests/default.nix
+++ b/tests/default.nix
@@ -298,7 +298,7 @@ in
 
   devShell-override = test
     (flakelight ./empty {
-      devShells.default = { mkShell }: mkShell { };
+      devShell = { mkShell }: mkShell { };
     })
     (f: f ? devShells.x86_64-linux.default);
 
@@ -314,6 +314,12 @@ in
       && (f ? devShells.x86_64-linux.shell1)
       && (f ? devShells.x86_64-linux.shell2));
 
+  devShells-override = test
+    (flakelight ./empty {
+      devShells.default = { mkShell }: mkShell { };
+    })
+    (f: f ? devShells.x86_64-linux.default);
+
   overlay = test
     (flakelight ./empty {
       overlay = final: prev: { testValue = "hello"; };