From 1240b904753eb6ae70b0d5ca3ed9e5d6428fab42 Mon Sep 17 00:00:00 2001
From: Nathan Sarang-Walters <nwalters512@gmail.com>
Date: Mon, 16 Dec 2024 15:09:31 -0800
Subject: [PATCH] Implement --check flag for fmt

---
 d2cli/fmt.go  | 24 +++++++++++++++++++++---
 d2cli/main.go |  7 ++++++-
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/d2cli/fmt.go b/d2cli/fmt.go
index a2b837158..61daf15ea 100644
--- a/d2cli/fmt.go
+++ b/d2cli/fmt.go
@@ -12,9 +12,10 @@ import (
 
 	"oss.terrastruct.com/d2/d2format"
 	"oss.terrastruct.com/d2/d2parser"
+	"oss.terrastruct.com/d2/lib/log"
 )
 
-func fmtCmd(ctx context.Context, ms *xmain.State) (err error) {
+func fmtCmd(ctx context.Context, ms *xmain.State, check bool) (err error) {
 	defer xdefer.Errorf(&err, "failed to fmt")
 
 	ms.Opts = xmain.NewOpts(ms.Env, ms.Opts.Flags.Args()[1:])
@@ -22,6 +23,8 @@ func fmtCmd(ctx context.Context, ms *xmain.State) (err error) {
 		return xmain.UsageErrorf("fmt must be passed at least one file to be formatted")
 	}
 
+	unformattedCount := 0
+
 	for _, inputPath := range ms.Opts.Args {
 		if inputPath != "-" {
 			inputPath = ms.AbsPath(inputPath)
@@ -43,10 +46,25 @@ func fmtCmd(ctx context.Context, ms *xmain.State) (err error) {
 
 		output := []byte(d2format.Format(m))
 		if !bytes.Equal(output, input) {
-			if err := ms.WritePath(inputPath, output); err != nil {
-				return err
+			if check {
+				unformattedCount += 1
+				log.Warn(ctx, inputPath)
+			} else {
+				if err := ms.WritePath(inputPath, output); err != nil {
+					return err
+				}
 			}
 		}
 	}
+
+	if unformattedCount > 0 {
+		pluralFiles := "file"
+		if unformattedCount > 1 {
+			pluralFiles = "files"
+		}
+
+		return xmain.ExitErrorf(1, "found %d unformatted %s. Run d2 fmt to fix.", unformattedCount, pluralFiles)
+	}
+
 	return nil
 }
diff --git a/d2cli/main.go b/d2cli/main.go
index eeefd6ae9..0240e9d7c 100644
--- a/d2cli/main.go
+++ b/d2cli/main.go
@@ -119,6 +119,11 @@ func Run(ctx context.Context, ms *xmain.State) (err error) {
 	fontBoldFlag := ms.Opts.String("D2_FONT_BOLD", "font-bold", "", "", "path to .ttf file to use for the bold font. If none provided, Source Sans Pro Bold is used.")
 	fontSemiboldFlag := ms.Opts.String("D2_FONT_SEMIBOLD", "font-semibold", "", "", "path to .ttf file to use for the semibold font. If none provided, Source Sans Pro Semibold is used.")
 
+	checkFlag, err := ms.Opts.Bool("D2_CHECK", "check", "", false, "check that the specified files are formatted correctly.")
+	if err != nil {
+		return err
+	}
+
 	plugins, err := d2plugin.ListPlugins(ctx)
 	if err != nil {
 		return err
@@ -153,7 +158,7 @@ func Run(ctx context.Context, ms *xmain.State) (err error) {
 			themesCmd(ctx, ms)
 			return nil
 		case "fmt":
-			return fmtCmd(ctx, ms)
+			return fmtCmd(ctx, ms, *checkFlag)
 		case "version":
 			if len(ms.Opts.Flags.Args()) > 1 {
 				return xmain.UsageErrorf("version subcommand accepts no arguments")