nix: move formatting checks to nix and enable in ci

.github/workflows/pipeline.yml

@@ -2,9 +2,6 @@ name: ci-pipeline
 
 on:
   push:
-    branches:
-      - staging
-      - main
   pull_request:
     branches:
       - staging
@@ -91,6 +88,24 @@ jobs:
 
       - uses: ./.github/actions/typecheck
 
+  nix-flake-check:
+    name: Check Nix flake
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          lfs: false
+
+      - uses: cachix/install-nix-action@v31
+        with:
+          nix_path: nixpkgs=channel:nixos-unstable
+
+      - name: Run nix flake check
+        run: |
+          nix flake check
+        shell: bash
+
 #  ci:
 #    needs: typecheck
 #    runs-on: ubuntu-latest

flake.lock

@@ -59,7 +59,8 @@
       "inputs": {
         "fenix": "fenix",
         "flake-utils": "flake-utils",
-        "nixpkgs": "nixpkgs"
+        "nixpkgs": "nixpkgs",
+        "treefmt-nix": "treefmt-nix"
       }
     },
     "rust-analyzer-src": {
@@ -93,6 +94,26 @@
         "repo": "default",
         "type": "github"
       }
+    },
+    "treefmt-nix": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1762938485,
+        "narHash": "sha256-AlEObg0syDl+Spi4LsZIBrjw+snSVU4T8MOeuZJUJjM=",
+        "owner": "numtide",
+        "repo": "treefmt-nix",
+        "rev": "5b4ee75aeefd1e2d5a1cc43cf6ba65eba75e83e4",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "treefmt-nix",
+        "type": "github"
+      }
     }
   },
   "root": "root",

flake.nix

@@ -9,6 +9,11 @@
       url = "github:nix-community/fenix";
       inputs.nixpkgs.follows = "nixpkgs";
     };
+    # Provides formatting infrastructure:
+    treefmt-nix = {
+      url = "github:numtide/treefmt-nix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
   };
 
   # TODO: figure out caching story
@@ -26,6 +31,7 @@
         "aarch64-darwin"
         "aarch64-linux"
       ];
+      fenixToolchain = system: inputs.fenix.packages.${system}.complete;
     in
     inputs.flake-utils.lib.eachSystem systems (
       system:
@@ -34,8 +40,19 @@
           inherit system;
           overlays = [ inputs.fenix.overlays.default ];
         };
+        treefmtEval = inputs.treefmt-nix.lib.evalModule pkgs {
+          projectRootFile = "flake.nix";
+          programs.ruff-format.enable = true;
+          programs.ruff-format.excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ];
+          programs.rustfmt.enable = true;
+          programs.rustfmt.package = (fenixToolchain system).rustfmt;
+          programs.nixpkgs-fmt.enable = true;
+        };
       in
       {
+        formatter = treefmtEval.config.build.wrapper;
+        checks.formatting = treefmtEval.config.build.check inputs.self;
+
         devShells.default = pkgs.mkShell {
           packages =
             with pkgs;
@@ -47,7 +64,7 @@
               basedpyright
 
               # RUST
-              (fenix.complete.withComponents [
+              ((fenixToolchain system).withComponents [
                 "cargo"
                 "rustc"
                 "clippy"

justfile

@@ -1,5 +1,5 @@
 fmt:
-    uv run ruff format src .mlx_typings
+    nix fmt
 
 lint:
     uv run ruff check --fix src