From 22f03b8f48a582bdc62945f871e51404b8fdfb49 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BAlio=20C=C3=A9sar=20Batista?= <julio@terrastruct.com>
Date: Thu, 13 Apr 2023 10:54:26 -0300
Subject: [PATCH] add exportExtension

---
 d2cli/export.go      | 30 ++++++++++++++++++
 d2cli/export_test.go | 73 ++++++++++++++++++++++++++++++++++++++++++++
 d2cli/main.go        | 15 ++++-----
 3 files changed, 111 insertions(+), 7 deletions(-)
 create mode 100644 d2cli/export.go
 create mode 100644 d2cli/export_test.go

diff --git a/d2cli/export.go b/d2cli/export.go
new file mode 100644
index 000000000..8cd2f1601
--- /dev/null
+++ b/d2cli/export.go
@@ -0,0 +1,30 @@
+package d2cli
+
+import "path/filepath"
+
+type exportExtension string
+
+var KNOWN_EXTENSIONS = []string{".svg", ".png", ".pptx", ".pdf"}
+
+func getExportExtension(outputPath string) exportExtension {
+	ext := filepath.Ext(outputPath)
+	for _, kext := range KNOWN_EXTENSIONS {
+		if kext == ext {
+			return exportExtension(ext)
+		}
+	}
+	// default is svg
+	return exportExtension(".svg")
+}
+
+func (ex exportExtension) supportsAnimation() bool {
+	return ex == ".svg"
+}
+
+func (ex exportExtension) requiresPngRenderer() bool {
+	return ex == ".png" || ex == ".pdf" || ex == ".pptx"
+}
+
+func (ex exportExtension) supportsDarkTheme() bool {
+	return ex == ".svg"
+}
diff --git a/d2cli/export_test.go b/d2cli/export_test.go
new file mode 100644
index 000000000..8518b1850
--- /dev/null
+++ b/d2cli/export_test.go
@@ -0,0 +1,73 @@
+package d2cli
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestOutputFormat(t *testing.T) {
+	type testCase struct {
+		outputPath        string
+		extension         exportExtension
+		supportsDarkTheme bool
+		supportsAnimation bool
+		requiresPngRender bool
+	}
+	testCases := []testCase{
+		{
+			outputPath:        "/out.svg",
+			extension:         ".svg",
+			supportsDarkTheme: true,
+			supportsAnimation: true,
+			requiresPngRender: false,
+		},
+		{
+			// assumes SVG by default
+			outputPath:        "/out",
+			extension:         ".svg",
+			supportsDarkTheme: true,
+			supportsAnimation: true,
+			requiresPngRender: false,
+		},
+		{
+			outputPath:        "-",
+			extension:         ".svg",
+			supportsDarkTheme: true,
+			supportsAnimation: true,
+			requiresPngRender: false,
+		},
+		{
+			outputPath:        "/out.png",
+			extension:         ".png",
+			supportsDarkTheme: false,
+			supportsAnimation: false,
+			requiresPngRender: true,
+		},
+		{
+			outputPath:        "/out.pptx",
+			extension:         ".pptx",
+			supportsDarkTheme: false,
+			supportsAnimation: false,
+			requiresPngRender: true,
+		},
+		{
+			outputPath:        "/out.pdf",
+			extension:         ".pdf",
+			supportsDarkTheme: false,
+			supportsAnimation: false,
+			requiresPngRender: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.outputPath, func(t *testing.T) {
+			extension := getExportExtension(tc.outputPath)
+			assert.Equal(t, tc.extension, extension)
+			assert.Equal(t, tc.supportsAnimation, extension.supportsAnimation())
+			assert.Equal(t, tc.supportsDarkTheme, extension.supportsDarkTheme())
+			assert.Equal(t, tc.requiresPngRender, extension.requiresPngRenderer())
+		})
+	}
+}
diff --git a/d2cli/main.go b/d2cli/main.go
index 58f867291..1ceaf7371 100644
--- a/d2cli/main.go
+++ b/d2cli/main.go
@@ -186,14 +186,13 @@ func Run(ctx context.Context, ms *xmain.State) (err error) {
 			inputPath = filepath.Join(inputPath, "index.d2")
 		}
 	}
+	outputFormat := getExportExtension(outputPath)
 	if outputPath != "-" {
 		outputPath = ms.AbsPath(outputPath)
-		if *animateIntervalFlag > 0 {
-			// Not checking for extension == "svg", because users may want to write SVG data to a non-svg-extension file
-			if filepath.Ext(outputPath) == ".png" || filepath.Ext(outputPath) == ".pdf" || filepath.Ext(outputPath) == ".pptx" {
-				return xmain.UsageErrorf("-animate-interval can only be used when exporting to SVG.\nYou provided: %s", filepath.Ext(outputPath))
-			}
+		if *animateIntervalFlag > 0 && !outputFormat.supportsAnimation() {
+			return xmain.UsageErrorf("-animate-interval can only be used when exporting to SVG.\nYou provided: %s", filepath.Ext(outputPath))
 		}
+
 	}
 
 	match := d2themescatalog.Find(*themeFlag)
@@ -236,12 +235,14 @@ func Run(ctx context.Context, ms *xmain.State) (err error) {
 	}
 	ms.Log.Debug.Printf("using layout plugin %s (%s)", *layoutFlag, plocation)
 
-	var pw png.Playwright
-	if filepath.Ext(outputPath) == ".png" || filepath.Ext(outputPath) == ".pdf" || filepath.Ext(outputPath) == ".pptx" {
+	if !outputFormat.supportsDarkTheme() {
 		if darkThemeFlag != nil {
 			ms.Log.Warn.Printf("--dark-theme cannot be used while exporting to another format other than .svg")
 			darkThemeFlag = nil
 		}
+	}
+	var pw png.Playwright
+	if outputFormat.requiresPngRenderer() {
 		pw, err = png.InitPlaywright()
 		if err != nil {
 			return err