diff --git a/.envrc b/.envrc
new file mode 100644
index 0000000..af0cc93
--- /dev/null
+++ b/.envrc
@@ -0,0 +1,2 @@
+#!/usr/bin/env bash
+use flake
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..8c9c964
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+.direnv
+.rumdl_cache
diff --git a/README.md b/README.md
index cdfd1da..abe202f 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,7 @@
 # nix-userstyles
 
-This module exports a nix function to generate userstyles for popular websites using your own color palette.
+This module exports a nix function to generate userstyles for popular
+websites using your own color palette.
 
 ![nix-userstyles](screenshot.png)
 
@@ -36,9 +37,13 @@ let
   ];
 in
 {
-  # no additional extensions needed, just add the userstyles to your firefox profile userContent
+  # no additional extensions needed, just add the userstyles
+  # to your firefox profile userContent
+  system = pkgs.stdenv.hostPlatform.system;
   programs.firefox.profiles.yourprofile.userContent = ''
-    ${builtins.readFile "${nix-userstyles.packages.${pkgs.stdenv.hostPlatform.system}.mkUserStyles palette userStyles}"}
+    ${builtins.readFile "${
+      nix-userstyles.packages.${system}.mkUserStyles palette userStyles
+    }"}
   '';
 }
 ```
@@ -57,7 +62,9 @@ let
     hasPrefix
     ;
 
-  palette = config.lib.stylix.colors |> filterAttrs (name: _: hasPrefix "base0" name);
+  palette =
+    config.lib.stylix.colors
+    |> filterAttrs (name: _: hasPrefix "base0" name);
   userStyles = [
     "brave-search"
     "bsky"
@@ -77,9 +84,13 @@ let
   ];
 in
 {
-  # no additional extensions needed, just add the userstyles to your firefox profile userContent
+  # no additional extensions needed, just add the userstyles
+  # to your firefox profile userContent
+  system = pkgs.stdenv.hostPlatform.system;
   programs.firefox.profiles.yourprofile.userContent = ''
-    ${builtins.readFile "${nix-userstyles.packages.${pkgs.stdenv.hostPlatform.system}.mkUserStyles palette userStyles}"}
+    ${builtins.readFile "${
+      nix-userstyles.packages.${system}.mkUserStyles palette userStyles
+    }"}
   '';
 }
 ```
diff --git a/flake.lock b/flake.lock
index 0dce935..9168634 100644
--- a/flake.lock
+++ b/flake.lock
@@ -44,6 +44,39 @@
         "url": "https://catppuccin.github.io/discord/dist/catppuccin-mocha.theme.css"
       }
     },
+    "flake-parts": {
+      "inputs": {
+        "nixpkgs-lib": "nixpkgs-lib"
+      },
+      "locked": {
+        "lastModified": 1769996383,
+        "narHash": "sha256-AnYjnFWgS49RlqX7LrC4uA+sCCDBj0Ry/WOJ5XWAsa0=",
+        "owner": "hercules-ci",
+        "repo": "flake-parts",
+        "rev": "57928607ea566b5db3ad13af0e57e921e6b12381",
+        "type": "github"
+      },
+      "original": {
+        "owner": "hercules-ci",
+        "repo": "flake-parts",
+        "type": "github"
+      }
+    },
+    "import-tree": {
+      "locked": {
+        "lastModified": 1763762820,
+        "narHash": "sha256-ZvYKbFib3AEwiNMLsejb/CWs/OL/srFQ8AogkebEPF0=",
+        "owner": "vic",
+        "repo": "import-tree",
+        "rev": "3c23749d8013ec6daa1d7255057590e9ca726646",
+        "type": "github"
+      },
+      "original": {
+        "owner": "vic",
+        "repo": "import-tree",
+        "type": "github"
+      }
+    },
     "nix-colors": {
       "inputs": {
         "base16-schemes": "base16-schemes",
@@ -80,19 +113,18 @@
         "type": "indirect"
       }
     },
-    "nixpkgs_2": {
+    "nixpkgs-lib": {
       "locked": {
-        "lastModified": 1770107345,
-        "narHash": "sha256-tbS0Ebx2PiA1FRW8mt8oejR0qMXmziJmPaU1d4kYY9g=",
-        "owner": "nixos",
-        "repo": "nixpkgs",
-        "rev": "4533d9293756b63904b7238acb84ac8fe4c8c2c4",
+        "lastModified": 1770515899,
+        "narHash": "sha256-hbmM5OSFCXIyoYvmZyQL9mjQ2mh/L1+2/4gf/BpXWNE=",
+        "owner": "nix-community",
+        "repo": "nixpkgs.lib",
+        "rev": "fd928847a8e03461e4b37699e6218539b610217d",
         "type": "github"
       },
       "original": {
-        "owner": "nixos",
-        "ref": "nixpkgs-unstable",
-        "repo": "nixpkgs",
+        "owner": "nix-community",
+        "repo": "nixpkgs.lib",
         "type": "github"
       }
     },
@@ -100,14 +132,34 @@
       "inputs": {
         "catppuccin-userstyles": "catppuccin-userstyles",
         "discord-userstyle": "discord-userstyle",
+        "flake-parts": "flake-parts",
+        "import-tree": "import-tree",
         "nix-colors": "nix-colors",
         "nixpkgs": "nixpkgs",
+        "systems": "systems",
         "treefmt-nix": "treefmt-nix"
       }
     },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    },
     "treefmt-nix": {
       "inputs": {
-        "nixpkgs": "nixpkgs_2"
+        "nixpkgs": [
+          "nixpkgs"
+        ]
       },
       "locked": {
         "lastModified": 1770228511,
diff --git a/flake.nix b/flake.nix
index d24c0f7..ac5d611 100644
--- a/flake.nix
+++ b/flake.nix
@@ -2,59 +2,39 @@
   inputs = {
     nixpkgs.url = "nixpkgs/nixpkgs-unstable";
 
-    nix-colors.url = "github:misterio77/nix-colors";
-    nix-colors.inputs.nixpkgs-lib.follows = "nixpkgs";
-
-    catppuccin-userstyles.url = "github:catppuccin/userstyles";
-    catppuccin-userstyles.flake = false;
-
-    discord-userstyle.url = "https://catppuccin.github.io/discord/dist/catppuccin-mocha.theme.css";
-    discord-userstyle.flake = false;
-
-    treefmt-nix.url = "github:numtide/treefmt-nix";
+    flake-parts.url = "github:hercules-ci/flake-parts";
+    import-tree.url = "github:vic/import-tree";
+    systems.url = "github:nix-systems/default";
+
+    nix-colors = {
+      url = "github:misterio77/nix-colors";
+      inputs.nixpkgs-lib.follows = "nixpkgs";
+    };
+
+    catppuccin-userstyles = {
+      url = "github:catppuccin/userstyles";
+      flake = false;
+    };
+
+    discord-userstyle = {
+      url = "https://catppuccin.github.io/discord/dist/catppuccin-mocha.theme.css";
+      flake = false;
+    };
+
+    treefmt-nix = {
+      url = "github:numtide/treefmt-nix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
   };
 
-  outputs = {
-    nixpkgs,
-    treefmt-nix,
+  outputs = inputs @ {
+    flake-parts,
+    import-tree,
+    systems,
     ...
-  } @ inputs: let
-    systems = [
-      "x86_64-linux"
-      "aarch64-linux"
-      "x86_64-darwin"
-      "aarch64-darwin"
-    ];
-    forAllSystems = f:
-      builtins.listToAttrs (
-        map (system: {
-          name = system;
-          value = f system;
-        })
-        systems
-      );
-
-    # Small tool to iterate over each systems
-    eachSystem = f: nixpkgs.lib.genAttrs systems (system: f nixpkgs.legacyPackages.${system});
-
-    # Eval the treefmt modules from ./treefmt.nix
-    treefmtEval = eachSystem (pkgs: treefmt-nix.lib.evalModule pkgs ./treefmt.nix);
-  in {
-    packages = forAllSystems (
-      system: let
-        pkgs = import nixpkgs {
-          inherit system;
-          inherit (inputs) nixpkgs;
-        };
-      in {
-        mkUserStyles = pkgs.callPackage ./lib/mkUserStyles.nix inputs;
-      }
-    );
-
-    formatter = eachSystem (pkgs: treefmtEval.${pkgs.stdenv.hostPlatform.system}.config.build.wrapper);
-    # for `nix flake check`
-    checks = eachSystem (pkgs: {
-      formatting = treefmtEval.${pkgs.stdenv.hostPlatform.system}.config.build.check inputs.self;
-    });
-  };
+  }:
+    flake-parts.lib.mkFlake {inherit inputs;} {
+      systems = import systems;
+      imports = [(import-tree ./parts)];
+    };
 }
diff --git a/parts/devshell.nix b/parts/devshell.nix
new file mode 100644
index 0000000..aaa9023
--- /dev/null
+++ b/parts/devshell.nix
@@ -0,0 +1,9 @@
+{
+  perSystem = {pkgs, ...}: let
+    inherit (pkgs) mkShell;
+  in {
+    devShells.default = mkShell {
+      packages = [pkgs.tokei];
+    };
+  };
+}
diff --git a/parts/packages.nix b/parts/packages.nix
new file mode 100644
index 0000000..da8e52c
--- /dev/null
+++ b/parts/packages.nix
@@ -0,0 +1,13 @@
+{inputs, ...}: {
+  perSystem = {pkgs, ...}: let
+    mkUserStyles = import ../lib/mkUserStyles.nix {
+      inherit pkgs;
+      inherit (pkgs) lib;
+      inherit (inputs) nix-colors catppuccin-userstyles discord-userstyle;
+    };
+  in {
+    packages = {
+      inherit mkUserStyles;
+    };
+  };
+}
diff --git a/parts/tests.nix b/parts/tests.nix
new file mode 100644
index 0000000..c4ebeb2
--- /dev/null
+++ b/parts/tests.nix
@@ -0,0 +1,39 @@
+{inputs, ...}: let
+  testUserStyles = [
+    "brave-search"
+    "bsky"
+    "chatgpt"
+    "cinny"
+    "claude"
+    "devdocs"
+    "discord"
+    "duckduckgo"
+    "github"
+    "google"
+    "hacker-news"
+    "lobste.rs"
+    "nixos-*"
+    "npm"
+    "ollama"
+    "perplexity"
+    "qwant"
+    "reddit"
+    "spotify-web"
+    "stack-overflow"
+    "telegram"
+    "whatsapp-web"
+    "wikipedia"
+    "youtube"
+  ];
+in {
+  perSystem = {pkgs, ...}: let
+    inherit (inputs.nix-colors.colorSchemes.dracula) palette;
+    mkUserStyles = import ../lib/mkUserStyles.nix {
+      inherit pkgs;
+      inherit (pkgs) lib;
+      inherit (inputs) nix-colors catppuccin-userstyles discord-userstyle;
+    };
+  in {
+    packages.test = mkUserStyles palette testUserStyles;
+  };
+}
diff --git a/parts/treefmt.nix b/parts/treefmt.nix
new file mode 100644
index 0000000..6e7b48f
--- /dev/null
+++ b/parts/treefmt.nix
@@ -0,0 +1,23 @@
+{inputs, ...}: {
+  perSystem = {pkgs, ...}: let
+    treefmtEval = inputs.treefmt-nix.lib.evalModule pkgs {
+      programs = {
+        alejandra.enable = true;
+        nixf-diagnose.enable = true;
+        deadnix.enable = true;
+        statix.enable = true;
+
+        biome.enable = true;
+
+        yamlfmt.enable = true;
+        yamllint.enable = true;
+
+        rumdl-check.enable = true;
+        rumdl-format.enable = true;
+      };
+    };
+  in {
+    formatter = treefmtEval.config.build.wrapper;
+    checks.formatting = treefmtEval.config.build.check inputs.self;
+  };
+}
diff --git a/test-mkUserStyles.nix b/test-mkUserStyles.nix
deleted file mode 100644
index e652737..0000000
--- a/test-mkUserStyles.nix
+++ /dev/null
@@ -1,31 +0,0 @@
-let
-  flake = builtins.getFlake (toString ./.);
-  nix-colors = builtins.getFlake "github:misterio77/nix-colors";
-  inherit (nix-colors.outputs.colorSchemes.dracula) palette;
-in
-  flake.outputs.packages.${builtins.currentSystem}.mkUserStyles palette [
-    "brave-search"
-    "bsky"
-    "chatgpt"
-    "cinny"
-    "claude"
-    "devdocs"
-    "discord"
-    "duckduckgo"
-    "github"
-    "google"
-    "hacker-news"
-    "lobste.rs"
-    "nixos-*"
-    "npm"
-    "ollama"
-    "perplexity"
-    "qwant"
-    "reddit"
-    "spotify-web"
-    "stack-overflow"
-    "telegram"
-    "whatsapp-web"
-    "wikipedia"
-    "youtube"
-  ]
diff --git a/treefmt.nix b/treefmt.nix
deleted file mode 100644
index 302a47f..0000000
--- a/treefmt.nix
+++ /dev/null
@@ -1,13 +0,0 @@
-_: {
-  programs = {
-    alejandra.enable = true;
-    nixf-diagnose.enable = true;
-    deadnix.enable = true;
-    statix.enable = true;
-
-    biome.enable = true;
-
-    yamlfmt.enable = true;
-    yamllint.enable = true;
-  };
-}