diff --git a/flake.nix b/flake.nix
index f3440c3..629588f 100644
--- a/flake.nix
+++ b/flake.nix
@@ -48,20 +48,11 @@
         systems = [ "x86_64-linux" "aarch64-linux" "x86_64-darwin" "aarch64-darwin" ];
 
         imports = [
+          inputs.treefmt-nix.flakeModule
+          ./shell.nix
           ./treefmt.nix
         ];
 
-        perSystem =
-          { config
-          , inputs'
-          , pkgs
-          , self'
-          , ...
-          }: {
-            devShells.default = pkgs.callPackage ./shell.nix {
-              inherit config;
-            };
-          };
         flake.nixosConfigurations =
           let
             inherit (inputs.nixpkgs.lib) nixosSystem;
diff --git a/shell.nix b/shell.nix
index 6b0e1af..bd90878 100644
--- a/shell.nix
+++ b/shell.nix
@@ -1,21 +1,22 @@
-{ config
-, pkgs
-}:
-
-with pkgs;
-mkShellNoCC {
-  buildInputs = [
-    jq
-    sops
-    ssh-to-age
-    (python3.withPackages (
-      p: [
-        p.deploykit
-        p.invoke
-        p.requests
-      ]
-    ))
-    rsync
-    config.treefmt.build.wrapper
-  ];
+{
+  perSystem = { config, pkgs, ... }: {
+    devShells = {
+      default = with pkgs; mkShellNoCC {
+        buildInputs = [
+          jq
+          sops
+          ssh-to-age
+          (python3.withPackages (
+            p: [
+              p.deploykit
+              p.invoke
+              p.requests
+            ]
+          ))
+          rsync
+          config.treefmt.build.wrapper
+        ];
+      };
+    };
+  };
 }
diff --git a/treefmt.nix b/treefmt.nix
index a8e6a60..bfa4a97 100644
--- a/treefmt.nix
+++ b/treefmt.nix
@@ -1,7 +1,4 @@
-{ inputs, ... }: {
-  imports = [
-    inputs.treefmt-nix.flakeModule
-  ];
+{
   perSystem = { pkgs, ... }: {
     treefmt = {
       # Used to find the project root