diff --git a/builtinModules/devShells.nix b/builtinModules/devShells.nix
index 0d2ad77..c42bbef 100644
--- a/builtinModules/devShells.nix
+++ b/builtinModules/devShells.nix
@@ -2,13 +2,15 @@
 # Copyright (C) 2023 Archit Gupta <archit@accelbread.com>
 # SPDX-License-Identifier: MIT
 
-{ config, lib, flakelight, genSystems, ... }:
+{ config, lib, flakelight, genSystems, moduleArgs, ... }:
 let
-  inherit (lib) filterAttrs mapAttrs mkDefault mkIf mkMerge mkOption;
+  inherit (builtins) attrNames hasAttr;
+  inherit (lib) all filterAttrs functionArgs mapAttrs mkDefault mkIf mkMerge
+    mkOption;
   inherit (lib.types) coercedTo functionTo lazyAttrsOf lines listOf nullOr
     package str submodule;
   inherit (flakelight) supportedSystem;
-  inherit (flakelight.types) optFunctionTo packageDef;
+  inherit (flakelight.types) function optFunctionTo packageDef;
 
   devShellModule.options = {
     inputsFrom = mkOption {
@@ -42,13 +44,17 @@ let
       default = null;
     };
   };
+
+  moduleFromFn = fn:
+    if all (a: hasAttr a moduleArgs) (attrNames (functionArgs fn))
+    then fn moduleArgs
+    else { overrideShell = fn; };
 in
 {
   options = {
     devShell = mkOption {
       default = null;
-      type = nullOr (coercedTo packageDef
-        (x: { overrideShell = x; })
+      type = nullOr (coercedTo function moduleFromFn
         (submodule devShellModule));
     };
 
diff --git a/builtinModules/nixDir.nix b/builtinModules/nixDir.nix
index ee6bdee..cf02b2f 100644
--- a/builtinModules/nixDir.nix
+++ b/builtinModules/nixDir.nix
@@ -4,7 +4,7 @@
 
 { config, src, lib, flakelight, moduleArgs, ... }:
 let
-  inherit (lib) mkOption mkIf mkMerge optionalAttrs;
+  inherit (lib) mkOption mkIf mkMerge;
   inherit (flakelight) autoImport autoImportArgs;
   inherit (flakelight.types) path;
 
@@ -25,7 +25,7 @@ in
       package = autoImport' "package";
       packages = autoImportArgs' "packages";
       overlays = autoImportArgs' "overlays";
-      devShell = autoImportArgs' "devShell";
+      devShell = autoImport' "devShell";
       devShells = autoImportArgs' "devShells";
       app = autoImport' "app";
       apps = autoImport' "apps";
@@ -53,16 +53,7 @@ in
       (mkIf (package != null) { inherit package; })
       (mkIf (packages != null) { inherit packages; })
       (mkIf (overlays != null) { inherit overlays; })
-      (mkIf (devShell != null) {
-        devShell = optionalAttrs (devShell ? inputsFrom)
-          { inherit (devShell) inputsFrom; }
-        // optionalAttrs (devShell ? packages)
-          { inherit (devShell) packages; }
-        // optionalAttrs (devShell ? shellHook)
-          { inherit (devShell) shellHook; }
-        // optionalAttrs (devShell ? env)
-          { inherit (devShell) env; };
-      })
+      (mkIf (devShell != null) { inherit devShell; })
       (mkIf (devShells != null) { inherit devShells; })
       (mkIf (app != null) { inherit app; })
       (mkIf (apps != null) { inherit apps; })