luciano/d2
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Use Playwright to export images to png

Browse source
This commit is contained in:
Bernard Xie 2022-11-16 17:41:53 -08:00
parent 6810cc8b26
commit c3a4c9aa1c
No known key found for this signature in database
GPG key ID: 3C3E0036CE0F892C
6 changed files with 187 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
cmd/d2/main.go
View file

@ -10,8 +10,7 @@ import (
"strings"
"time"
_ "embed"
"github.com/playwright-community/playwright-go"
"github.com/spf13/pflag"
"oss.terrastruct.com/d2"
@ -20,6 +19,7 @@ import (
"oss.terrastruct.com/d2/d2renderers/textmeasure"
"oss.terrastruct.com/d2/d2themes"
"oss.terrastruct.com/d2/d2themes/d2themescatalog"
"oss.terrastruct.com/d2/lib/png"
"oss.terrastruct.com/d2/lib/version"
"oss.terrastruct.com/d2/lib/xmain"
)
@ -113,12 +113,28 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
}
ms.Log.Debug.Printf("using layout plugin %s (%s)", envD2Layout, pluginLocation)
var pw *playwright.Playwright
var browser playwright.Browser
if filepath.Ext(outputPath) == ".png" {
pw, browser, err = png.InitPlaywright()
if err != nil {
return err
}
}
defer func() error {
err = png.Cleanup(pw, browser)
if err != nil {
return err
}
return nil
}()
if *watchFlag {
if inputPath == "-" {
return xmain.UsageErrorf("-w[atch] cannot be combined with reading input from stdin")
}
ms.Env.Setenv("LOG_TIMESTAMPS", "1")
w, err := newWatcher(ctx, ms, plugin, inputPath, outputPath)
w, err := newWatcher(ctx, ms, plugin, inputPath, outputPath, pw, browser)
if err != nil {
return err
}
@ -132,15 +148,16 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
_ = 343
}
_, err = compile(ctx, ms, plugin, inputPath, outputPath)
_, err = compile(ctx, ms, plugin, inputPath, outputPath, browser)
if err != nil {
return err
}
ms.Log.Success.Printf("successfully compiled %v to %v", inputPath, outputPath)
return nil
}
func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, inputPath, outputPath string) ([]byte, error) {
func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, inputPath, outputPath string, browser playwright.Browser) ([]byte, error) {
input, err := ms.ReadPath(inputPath)
if err != nil {
return nil, err
@ -165,12 +182,19 @@ func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, input
if err != nil {
return nil, err
}
svg, err = plugin.PostProcess(ctx, svg)
outputImage, err := plugin.PostProcess(ctx, svg)
if err != nil {
return nil, err
}
err = ms.WritePath(outputPath, svg)
if filepath.Ext(outputPath) == ".png" {
outputImage, err = png.ExportPNG(browser, svg)
if err != nil {
return nil, err
}
}
err = ms.WritePath(outputPath, outputImage)
if err != nil {
return nil, err
}

19
cmd/d2/watch.go
View file

@ -16,6 +16,7 @@ import (
"time"
"github.com/fsnotify/fsnotify"
"github.com/playwright-community/playwright-go"
"nhooyr.io/websocket"
"nhooyr.io/websocket/wsjson"
@ -61,6 +62,9 @@ type watcher struct {
resMu sync.Mutex
res *compileResult
browser playwright.Browser
pw *playwright.Playwright
}
type compileResult struct {
@ -68,7 +72,7 @@ type compileResult struct {
SVG string `json:"svg"`
}
func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plugin, inputPath, outputPath string) (*watcher, error) {
func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plugin, inputPath, outputPath string, pw *playwright.Playwright, browser playwright.Browser) (*watcher, error) {
ctx, cancel := context.WithCancel(ctx)
w := &watcher{
@ -83,6 +87,8 @@ func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plug
compileCh: make(chan struct{}, 1),
wsclients: make(map[*wsclient]struct{}),
browser: browser,
pw: pw,
}
err := w.init()
if err != nil {
@ -143,6 +149,7 @@ func (w *watcher) run() error {
func (w *watcher) close() {
w.wsclientsMu.Lock()
if w.closing {
w.wsclientsMu.Unlock()
return
@ -325,7 +332,15 @@ func (w *watcher) compileLoop(ctx context.Context) error {
recompiledPrefix = "re"
}
b, err := compile(ctx, w.ms, w.layoutPlugin, w.inputPath, w.outputPath)
if !w.browser.IsConnected() {
newBrowser, err := w.pw.Chromium.Launch()
if err != nil {
return err
}
w.browser = newBrowser
}
b, err := compile(ctx, w.ms, w.layoutPlugin, w.inputPath, w.outputPath, w.browser)
if err != nil {
err = fmt.Errorf("failed to %scompile: %w", recompiledPrefix, err)
w.ms.Log.Error.Print(err)

5
go.mod
View file

@ -33,19 +33,21 @@ require (
require (
cloud.google.com/go/compute v1.7.0 // indirect
github.com/andybalholm/cascadia v1.3.1 // indirect
github.com/danwakefield/fnmatch v0.0.0-20160403171240-cbb64ac3d964 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/dlclark/regexp2 v1.4.0 // indirect
github.com/fatih/color v1.13.0 // indirect
github.com/gin-gonic/gin v1.7.7 // indirect
github.com/go-playground/validator/v10 v10.10.0 // indirect
github.com/go-stack/stack v1.8.1 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/gorilla/websocket v1.4.2 // indirect
github.com/json-iterator/go v1.1.12 // indirect
github.com/klauspost/compress v1.13.6 // indirect
github.com/kr/pretty v0.3.0 // indirect
github.com/mattn/go-colorable v0.1.12 // indirect
github.com/mattn/go-isatty v0.0.14 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/playwright-community/playwright-go v0.2000.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/rogpeppe/go-internal v1.8.0 // indirect
github.com/ugorji/go/codec v1.2.6 // indirect
@ -57,6 +59,7 @@ require (
google.golang.org/genproto v0.0.0-20220822174746-9e6da59bd2fc // indirect
google.golang.org/protobuf v1.28.1 // indirect
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
gopkg.in/square/go-jose.v2 v2.6.0 // indirect
gopkg.in/yaml.v2 v2.4.0 // indirect
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
)

9
go.sum
View file

@ -88,6 +88,8 @@ github.com/cncf/xds/go v0.0.0-20211001041855-01bcc9b48dfe/go.mod h1:eXthEFrGJvWH
github.com/cncf/xds/go v0.0.0-20211011173535-cb28da3451f1/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/creack/pty v1.1.18 h1:n56/Zwd5o6whRC5PMGretI4IdRLlmBXYNjScPaBgsbY=
github.com/danwakefield/fnmatch v0.0.0-20160403171240-cbb64ac3d964 h1:y5HC9v93H5EPKqaS1UYVg1uYah5Xf51mBfIoWehClUQ=
github.com/danwakefield/fnmatch v0.0.0-20160403171240-cbb64ac3d964/go.mod h1:Xd9hchkHSWYkEqJwUGisez3G1QY8Ryz0sdWrLPMGjLk=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@ -130,6 +132,8 @@ github.com/go-playground/validator/v10 v10.2.0/go.mod h1:uOYAAleCW8F/7oMFd6aG0GO
github.com/go-playground/validator/v10 v10.4.1/go.mod h1:nlOn6nFhuKACm19sB/8EGNn9GlaMV7XkbRSipzJ0Ii4=
github.com/go-playground/validator/v10 v10.10.0 h1:I7mrTYv78z8k8VXa/qJlOlEXn/nBh+BF8dHX5nt/dr0=
github.com/go-playground/validator/v10 v10.10.0/go.mod h1:74x4gJWsvQexRdW8Pn3dXSGrTK4nAUsbPlLADvpJkos=
github.com/go-stack/stack v1.8.1 h1:ntEHSVwIt7PNXNpgPmVfMrNhLtgjlmnZha2kOpuRiDw=
github.com/go-stack/stack v1.8.1/go.mod h1:dcoOX6HbPZSZptuspn9bctJ+N/CnF5gGygcUP3XYfe4=
github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee h1:s+21KNqlpePfkah2I+gwHF8xmJWRjooY+5248k6m4A0=
github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo=
github.com/gobwas/pool v0.2.0 h1:QEmUOlnSjWtnpRGHF3SauEiOsy82Cup83Vf2LcMlnc8=
@ -223,6 +227,7 @@ github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/ad
github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
github.com/h2non/filetype v1.1.1/go.mod h1:319b3zT68BvV+WRj7cwy856M2ehB3HqNOt6sy1HndBY=
github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
@ -266,6 +271,8 @@ github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjY
github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 h1:KoWmjvw+nsYOo29YJK9vDA65RGE3NrOnUtO7a+RF9HU=
github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8/go.mod h1:HKlIX3XHQyzLZPlr7++PzdhaXEj94dEiJgZDTsxEqUI=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
github.com/playwright-community/playwright-go v0.2000.1 h1:2JViSHpJQ/UL/PO1Gg6gXV5IcXAAsoBJ3KG9L3wKXto=
github.com/playwright-community/playwright-go v0.2000.1/go.mod h1:1y9cM9b9dVHnuRWzED1KLM7FtbwTJC8ibDjI6MNqewU=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
@ -751,6 +758,8 @@ gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
gopkg.in/square/go-jose.v2 v2.6.0 h1:NGk74WTnPKBNUhNzQX7PYcTLUjoq7mzKk2OKbvwk2iI=
gopkg.in/square/go-jose.v2 v2.6.0/go.mod h1:M9dMgbHiYLoDGQrXy7OpJDJWiKiU//h+vD76mk0e1AI=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.2.3/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

23
lib/png/generate_png.js Normal file
View file

@ -0,0 +1,23 @@
async (imgString) => {
const exportMaxWidth = 4000;
const tempImg = new Image();
const loadImage = () => {
return new Promise((resolve, reject) => {
tempImg.onload = (event) => resolve(event.currentTarget);
tempImg.onerror = () => {
reject("error loading string as an image");
};
tempImg.src = imgString;
});
};
const img = await loadImage();
const canvas = document.createElement("canvas");
canvas.width = Math.min(img.width, exportMaxWidth);
canvas.height = (canvas.width * img.height) / img.width;
const ctx = canvas.getContext("2d");
if (!ctx) {
return "";
}
ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
return canvas.toDataURL("image/png");
}

103
lib/png/png.go Normal file
View file

@ -0,0 +1,103 @@
package png
import (
"bytes"
"encoding/base64"
"errors"
"fmt"
"os"
"os/exec"
"strings"
_ "embed"
"github.com/playwright-community/playwright-go"
)
func InitPlaywright() (*playwright.Playwright, playwright.Browser, error) {
// check if playwright driver/browsers are installed and up to date
// https://github.com/playwright-community/playwright-go/blob/8e8f670b5fa7ba5365ae4bfc123fea4aac359763/run.go#L64.
driver, err := playwright.NewDriver(&playwright.RunOptions{})
if err != nil {
return nil, nil, err
}
if _, err := os.Stat(driver.DriverBinaryLocation); errors.Is(err, os.ErrNotExist) {
err = playwright.Install()
if err != nil {
return nil, nil, err
}
} else if err == nil {
cmd := exec.Command(driver.DriverBinaryLocation, "--version")
output, err := cmd.Output()
if err != nil || !bytes.Contains(output, []byte(driver.Version)) {
err = playwright.Install()
if err != nil {
return nil, nil, err
}
}
}
pw, err := playwright.Run()
if err != nil {
return nil, nil, err
}
browser, err := pw.Chromium.Launch()
if err != nil {
return nil, nil, err
}
return pw, browser, nil
}
//go:embed generate_png.js
var genPNGScript string
func ExportPNG(browser playwright.Browser, svg []byte) (outputImage []byte, err error) {
var page playwright.Page
defer func() error {
err = page.Close()
if err != nil {
return err
}
return nil
}()
if browser == nil {
return nil, fmt.Errorf("browser is not initialized for png export")
}
page, err = browser.NewPage()
if err != nil {
return nil, err
}
encodedSVG := base64.StdEncoding.EncodeToString(svg)
pngInterface, err := page.Evaluate(genPNGScript, "data:image/svg+xml;charset=utf-8;base64,"+encodedSVG)
if err != nil {
return nil, err
}
pngString := fmt.Sprintf("%v", pngInterface)
pngPrefix := "data:image/png;base64,"
if !strings.HasPrefix(pngString, pngPrefix) {
return nil, fmt.Errorf("playwright export generated invalid png")
}
splicedPNGString := pngString[len(pngPrefix):]
outputImage, err = base64.StdEncoding.DecodeString(splicedPNGString)
if err != nil {
return nil, err
}
return outputImage, nil
}
func Cleanup(pw *playwright.Playwright, browser playwright.Browser) (err error) {
if browser != nil {
if err = browser.Close(); err != nil {
return err
}
}
if pw != nil {
if err = pw.Stop(); err != nil {
return err
}
}
return nil
}