2 files changed, 22 insertions, 11 deletions

diff --git a/lib/mk-dev-shell.nix b/lib/mk-dev-shell.nix
index 8078f35..c6b36b7 100644
--- a/lib/mk-dev-shell.nix
+++ b/lib/mk-dev-shell.nix
@@ -1,32 +1,41 @@
 { pkgs
 , tools ? []
 , envVars ? {}
+, shell ? null
 }:
 
+with pkgs.lib;
+
 let
-  exportEnvVar = k: v: with pkgs.lib; ''
+  exportEnvVar = k: v: ''
     export ${escapeShellArg k}=${escapeShellArg v}
     echo ${escapeShellArg k}=${escapeShellArg v}
   '';
-  exportedEnvVars = pkgs.lib.mapAttrsToList exportEnvVar envVars;
-
-in pkgs.mkShell {
-  buildInputs = tools;
 
-  # TODO: handle case with no env variable set
-  # TODO: handle case with no package in tools
-  # TODO: allow custom shell
-  shellHook = ''
+  envVarsHook = envVars: ''
     echo -e "\nDEVSHELL ENVIRONMENT VARIABLES:"
-    ${pkgs.lib.concatStringsSep "\n" exportedEnvVars}
+    ${pkgs.lib.concatStringsSep "\n" (
+      pkgs.lib.mapAttrsToList exportEnvVar envVars
+    )}
+  '';
 
+  listToolsHook = tools: ''
     echo -e "\nDEVSHELL COMMANDS:"
     ls "${pkgs.symlinkJoin { name = "env"; paths = tools; }}/bin"
+  '';
 
-    # Use the default user shell instead of Bash
+  # Use the default user shell instead of Bash
+  startUserShellHook = ''
     $(${pkgs.finger_bsd}/bin/finger $USER \
       | ${pkgs.gnugrep}/bin/grep -oP 'Shell: \K.*')
+  '';
 
+in pkgs.mkShell {
+  buildInputs = tools;
+  shellHook = ''
+    ${optionalString (envVars != {}) (envVarsHook envVars)}
+    ${optionalString (tools != []) (listToolsHook tools)}
+    ${if (shell != null) then shell else startUserShellHook}
     exit $?
   '';
 }
diff --git a/readme.md b/readme.md
index 88c4d4a..e4ba8dd 100644
--- a/readme.md
+++ b/readme.md
@@ -48,6 +48,8 @@ Functions documentation provided as comments below.
         PGDATABASE = "app";
         DATABASE_URL = "postgresql:///${PGDATABASE}?host=${PGHOST}";
       };
+
+      shell = null;
     };
 
   });