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

create Playwright struct, use BrowserContext

Browse source
This commit is contained in:
Bernard Xie 2022-11-17 11:25:07 -08:00
parent e5e34362fe
commit 6ef4246565
No known key found for this signature in database
GPG key ID: 3C3E0036CE0F892C
3 changed files with 92 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
cmd/d2/main.go
View file

@ -113,15 +113,14 @@ 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
var pw png.Playwright
if filepath.Ext(outputPath) == ".png" {
pw, browser, err = png.InitPlaywright()
pw, err = png.InitPlaywright()
if err != nil {
return err
}
defer func() error {
err = png.Cleanup(pw, browser)
err = pw.Cleanup(*watchFlag)
if err != nil {
return err
}
@ -134,7 +133,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
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, pw, browser)
w, err := newWatcher(ctx, ms, plugin, inputPath, outputPath, pw)
if err != nil {
return err
}
@ -148,7 +147,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
_ = 343
}
_, err = compile(ctx, ms, plugin, inputPath, outputPath, browser)
_, err = compile(ctx, ms, plugin, inputPath, outputPath, pw.Page)
if err != nil {
return err
}
@ -157,7 +156,7 @@ func run(ctx context.Context, ms *xmain.State) (err error) {
return nil
}
func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, inputPath, outputPath string, browser playwright.Browser) ([]byte, error) {
func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, inputPath, outputPath string, page playwright.Page) ([]byte, error) {
input, err := ms.ReadPath(inputPath)
if err != nil {
return nil, err
@ -188,7 +187,7 @@ func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, input
}
if filepath.Ext(outputPath) == ".png" {
outputImage, err = png.ExportPNG(browser, svg)
outputImage, err = png.ExportPNG(ms, page, svg)
if err != nil {
return nil, err
}

17
cmd/d2/watch.go
View file

@ -16,11 +16,11 @@ import (
"time"
"github.com/fsnotify/fsnotify"
"github.com/playwright-community/playwright-go"
"nhooyr.io/websocket"
"nhooyr.io/websocket/wsjson"
"oss.terrastruct.com/d2/d2plugin"
"oss.terrastruct.com/d2/lib/png"
"oss.terrastruct.com/d2/lib/xbrowser"
"oss.terrastruct.com/d2/lib/xhttp"
"oss.terrastruct.com/d2/lib/xmain"
@ -63,8 +63,7 @@ type watcher struct {
resMu sync.Mutex
res *compileResult
browser playwright.Browser
pw *playwright.Playwright
pw png.Playwright
}
type compileResult struct {
@ -72,7 +71,7 @@ type compileResult struct {
SVG string `json:"svg"`
}
func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plugin, inputPath, outputPath string, pw *playwright.Playwright, browser playwright.Browser) (*watcher, error) {
func newWatcher(ctx context.Context, ms *xmain.State, layoutPlugin d2plugin.Plugin, inputPath, outputPath string, pw png.Playwright) (*watcher, error) {
ctx, cancel := context.WithCancel(ctx)
w := &watcher{
@ -87,7 +86,6 @@ 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()
@ -332,15 +330,16 @@ func (w *watcher) compileLoop(ctx context.Context) error {
recompiledPrefix = "re"
}
if !w.browser.IsConnected() {
newBrowser, err := w.pw.Chromium.Launch()
if filepath.Ext(w.outputPath) == ".png" && !w.pw.Browser.IsConnected() {
newPW, err := w.pw.RestartBrowser()
if err != nil {
w.ms.Log.Error.Printf("failed to refresh Playwright browser")
return err
}
w.browser = newBrowser
w.pw = newPW
}
b, err := compile(ctx, w.ms, w.layoutPlugin, w.inputPath, w.outputPath, w.browser)
b, err := compile(ctx, w.ms, w.layoutPlugin, w.inputPath, w.outputPath, w.pw.Page)
if err != nil {
err = fmt.Errorf("failed to %scompile: %w", recompiledPrefix, err)
w.ms.Log.Error.Print(err)

120
lib/png/png.go
View file

@ -12,19 +12,69 @@ import (
_ "embed"
"github.com/playwright-community/playwright-go"
"oss.terrastruct.com/d2/lib/xmain"
)
func InitPlaywright() (*playwright.Playwright, playwright.Browser, error) {
type Playwright struct {
PW *playwright.Playwright
Browser playwright.Browser
BrowserContext playwright.BrowserContext
Page playwright.Page
}
func (pw *Playwright) RestartBrowser() (newPW Playwright, err error) {
if err = pw.BrowserContext.Close(); err != nil {
return Playwright{}, err
}
if err = pw.Browser.Close(); err != nil {
return Playwright{}, err
}
browser, err := pw.PW.Chromium.Launch()
if err != nil {
return Playwright{}, err
}
context, err := browser.NewContext()
if err != nil {
return Playwright{}, err
}
page, err := context.NewPage()
if err != nil {
return Playwright{}, err
}
return Playwright{
PW: pw.PW,
Browser: browser,
BrowserContext: context,
Page: page,
}, nil
}
func (pw *Playwright) Cleanup(isWatch bool) (err error) {
if !isWatch {
if err = pw.BrowserContext.Close(); err != nil {
return err
}
}
if err = pw.Browser.Close(); err != nil {
return err
}
if err = pw.PW.Stop(); err != nil {
return err
}
return nil
}
func InitPlaywright() (Playwright, 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
return Playwright{}, err
}
if _, err := os.Stat(driver.DriverBinaryLocation); errors.Is(err, os.ErrNotExist) {
err = playwright.Install()
if err != nil {
return nil, nil, err
return Playwright{}, err
}
} else if err == nil {
cmd := exec.Command(driver.DriverBinaryLocation, "--version")
@ -32,42 +82,44 @@ func InitPlaywright() (*playwright.Playwright, playwright.Browser, error) {
if err != nil || !bytes.Contains(output, []byte(driver.Version)) {
err = playwright.Install()
if err != nil {
return nil, nil, err
return Playwright{}, err
}
}
}
pw, err := playwright.Run()
if err != nil {
return nil, nil, err
return Playwright{}, err
}
browser, err := pw.Chromium.Launch()
if err != nil {
return nil, nil, err
return Playwright{}, err
}
return pw, browser, nil
context, err := browser.NewContext()
if err != nil {
return Playwright{}, err
}
page, err := context.NewPage()
if err != nil {
return Playwright{}, err
}
return Playwright{
PW: pw,
Browser: browser,
BrowserContext: context,
Page: page,
}, 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
}()
func ExportPNG(ms *xmain.State, page playwright.Page, svg []byte) (outputImage []byte, err error) {
if page == nil {
ms.Log.Error.Printf("Playwright was not initialized properly for PNG export")
return nil, fmt.Errorf("Playwright page is not initialized for png export")
}
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 {
@ -77,27 +129,9 @@ func ExportPNG(browser playwright.Browser, svg []byte) (outputImage []byte, err
pngString := fmt.Sprintf("%v", pngInterface)
pngPrefix := "data:image/png;base64,"
if !strings.HasPrefix(pngString, pngPrefix) {
ms.Log.Error.Printf("failed to convert D2 file to PNG")
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
return base64.StdEncoding.DecodeString(splicedPNGString)
}