From d3931c4481751ece09bb0f7ca5fbab300a885740 Mon Sep 17 00:00:00 2001
From: jolheiser <john.olheiser@gmail.com>
Date: Wed, 19 Jun 2024 20:31:46 -0500
Subject: [PATCH] feat: add fmt

Signed-off-by: jolheiser <john.olheiser@gmail.com>
---
 nixfig.go      | 43 +++++++++++++++++++++++++++++++++++++++++--
 nixfig_test.go | 20 ++++++++++++++++++++
 2 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/nixfig.go b/nixfig.go
index b91d349..d89e98b 100644
--- a/nixfig.go
+++ b/nixfig.go
@@ -7,12 +7,14 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"strings"
 )
 
 var (
-	// Nix is the command to call for nix
-	Nix            string
+	Nix            string   // Nix is the command to call for nix
+	Fmt            []string // Fmt is the command (and args) to call to format the nix output
 	ErrNixNotFound = errors.New("nix was not found or set. You can set it either with `nixfig.Nix` or the `NIXFIG_NIX` environment variable")
+	ErrFmtNotFound = errors.New("nix formatter was not found or set. You can set it either with `nixfig.Fmt` or the `NIXFIG_FMT` environment variable")
 )
 
 func init() {
@@ -21,6 +23,17 @@ func init() {
 	if envPath, ok := os.LookupEnv("NIXFIG_NIX"); ok {
 		Nix = envPath
 	}
+
+	for _, formatter := range []string{"alejandra", "nixfmt-rfc-style", "nixfmt"} {
+		fmtPath, _ := exec.LookPath(formatter)
+		if fmtPath != "" {
+			Fmt = []string{fmtPath, "--quiet"}
+			break
+		}
+	}
+	if envPath, ok := os.LookupEnv("NIXFIG_FMT"); ok {
+		Fmt = strings.Split(envPath, " ")
+	}
 }
 
 // Unmarshal unmarshals a nix expression as JSON into a struct
@@ -70,3 +83,29 @@ func Marshal(v any) ([]byte, error) {
 
 	return stdout.Bytes(), nil
 }
+
+// MarshalFormat marshals a struct into a nix expression and formats it with Fmt
+func MarshalFormat(v any) ([]byte, error) {
+	if Fmt == nil {
+		return nil, ErrFmtNotFound
+	}
+
+	data, err := Marshal(v)
+	if err != nil {
+		return nil, err
+	}
+
+	var stdout, stderr bytes.Buffer
+	cmd := exec.Command(Fmt[0], Fmt[1:]...)
+	cmd.Stdin = bytes.NewBuffer(data)
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+	if err := cmd.Run(); err != nil {
+		return nil, fmt.Errorf("could not run %v: %w", Fmt, err)
+	}
+	if stderr.Len() > 0 {
+		return nil, errors.New(stderr.String())
+	}
+
+	return stdout.Bytes(), nil
+}
diff --git a/nixfig_test.go b/nixfig_test.go
index 1d70aee..bfde939 100644
--- a/nixfig_test.go
+++ b/nixfig_test.go
@@ -50,6 +50,16 @@ func TestNixNotFound(t *testing.T) {
 	Nix = oldNix
 }
 
+func TestFmtNotFound(t *testing.T) {
+	oldFmt := Fmt
+	Fmt = nil
+
+	_, err := MarshalFormat(nil)
+	assert.IsError(t, err, ErrFmtNotFound)
+
+	Fmt = oldFmt
+}
+
 func TestUnmarshal(t *testing.T) {
 	data, err := os.ReadFile("testdata/config.nix")
 	assert.NoError(t, err)
@@ -69,3 +79,13 @@ func TestMarshal(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, testCfg, cfg)
 }
+
+func TestMarshalFormat(t *testing.T) {
+	data, err := MarshalFormat(testCfg)
+	assert.NoError(t, err)
+
+	var cfg Config
+	err = Unmarshal(data, &cfg)
+	assert.NoError(t, err)
+	assert.Equal(t, testCfg, cfg)
+}