From 0f95288b5ca16a5664876789de0063090ad535d6 Mon Sep 17 00:00:00 2001
From: Bernard Xie <bernard@terrastruct.com>
Date: Fri, 18 Nov 2022 13:29:30 -0800
Subject: [PATCH] PR comments

---
 cmd/d2/main.go  |  2 +-
 cmd/d2/watch.go |  5 +++++
 lib/png/png.go  | 38 +++++++++++++++++++-------------------
 3 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/cmd/d2/main.go b/cmd/d2/main.go
index 7aeecf923..e8108cb74 100644
--- a/cmd/d2/main.go
+++ b/cmd/d2/main.go
@@ -186,7 +186,7 @@ func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, input
 	}
 
 	if filepath.Ext(outputPath) == ".png" {
-		outputImage, err = png.ExportPNG(ms, page, outputImage)
+		outputImage, err = png.ConvertSVG(ms, page, outputImage)
 		if err != nil {
 			return nil, err
 		}
diff --git a/cmd/d2/watch.go b/cmd/d2/watch.go
index 98a05bbdd..8207641de 100644
--- a/cmd/d2/watch.go
+++ b/cmd/d2/watch.go
@@ -332,6 +332,11 @@ func (w *watcher) compileLoop(ctx context.Context) error {
 		if filepath.Ext(w.outputPath) == ".png" && !w.pw.Browser.IsConnected() {
 			newPW, err := w.pw.RestartBrowser()
 			if err != nil {
+				broadcastErr := fmt.Errorf("issue encountered with PNG exporter, stopping")
+				w.ms.Log.Error.Print(broadcastErr)
+				w.broadcast(&compileResult{
+					Err: broadcastErr.Error(),
+				})
 				return fmt.Errorf("Playwright could not be restarted: %w", err)
 			}
 			w.pw = newPW
diff --git a/lib/png/png.go b/lib/png/png.go
index 90b947800..9a8d0d888 100644
--- a/lib/png/png.go
+++ b/lib/png/png.go
@@ -64,25 +64,28 @@ func InitPlaywright() (Playwright, error) {
 	if err != nil {
 		return Playwright{}, err
 	}
-	if _, err := os.Stat(driver.DriverBinaryLocation); os.IsNotExist(err) {
-		err = playwright.Install()
-		if err != nil {
-			return Playwright{}, err
-		}
-	} else if err == nil {
-		cmd := exec.Command(driver.DriverBinaryLocation, "--version")
-		output, err := cmd.Output()
-		if err != nil {
-			return Playwright{}, fmt.Errorf("error getting Playwright version: %w\nplease report this issue here: https://github.com/terrastruct/d2/issues/new", err)
-		}
-		if !bytes.Contains(output, []byte(driver.Version)) {
+	_, err = os.Stat(driver.DriverBinaryLocation)
+	if err != nil {
+		if os.IsNotExist(err) {
 			err = playwright.Install()
 			if err != nil {
 				return Playwright{}, err
 			}
+		} else {
+			return Playwright{}, fmt.Errorf("could not access Playwright binary location: %v\nerror: %w\nplease report this issue here: https://github.com/terrastruct/d2/issues/new", driver.DriverBinaryLocation, err)
+		}
+	}
+
+	cmd := exec.Command(driver.DriverBinaryLocation, "--version")
+	output, err := cmd.Output()
+	if err != nil {
+		return Playwright{}, fmt.Errorf("error getting Playwright version: %w\nplease report this issue here: https://github.com/terrastruct/d2/issues/new", err)
+	}
+	if !bytes.Contains(output, []byte(driver.Version)) {
+		err = playwright.Install()
+		if err != nil {
+			return Playwright{}, err
 		}
-	} else {
-		return Playwright{}, fmt.Errorf("could not access Playwright binary location: %v\nerror: %w\nplease report this issue here: https://github.com/terrastruct/d2/issues/new", driver.DriverBinaryLocation, err)
 	}
 
 	pw, err := playwright.Run()
@@ -95,11 +98,9 @@ func InitPlaywright() (Playwright, error) {
 //go:embed generate_png.js
 var genPNGScript string
 
-func ExportPNG(ms *xmain.State, page playwright.Page, svg []byte) (outputImage []byte, err error) {
-	if page == nil {
-		return nil, fmt.Errorf("Playwright was not initialized properly\nplease report this issue here: https://github.com/terrastruct/d2/issues/new")
-	}
+const pngPrefix = "data:image/png;base64,"
 
+func ConvertSVG(ms *xmain.State, page playwright.Page, svg []byte) (outputImage []byte, err error) {
 	encodedSVG := base64.StdEncoding.EncodeToString(svg)
 	pngInterface, err := page.Evaluate(genPNGScript, "data:image/svg+xml;charset=utf-8;base64,"+encodedSVG)
 	if err != nil {
@@ -107,7 +108,6 @@ func ExportPNG(ms *xmain.State, page playwright.Page, svg []byte) (outputImage [
 	}
 
 	pngString := fmt.Sprintf("%v", pngInterface)
-	pngPrefix := "data:image/png;base64,"
 	if !strings.HasPrefix(pngString, pngPrefix) {
 		if len(pngString) > 50 {
 			pngString = pngString[0:50] + "..."