png.go

Documentation: oss.terrastruct.com/d2/lib/png

     1  package png
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/base64"
     6  	"fmt"
     7  	"strings"
     8  
     9  	_ "embed"
    10  
    11  	exif "github.com/dsoprea/go-exif/v3"
    12  	exifcommon "github.com/dsoprea/go-exif/v3/common"
    13  	pngstruct "github.com/dsoprea/go-png-image-structure/v2"
    14  	"github.com/playwright-community/playwright-go"
    15  
    16  	"oss.terrastruct.com/d2/lib/version"
    17  )
    18  
    19  // ConvertSVG scales the image by 2x
    20  const SCALE = 2.
    21  
    22  type Playwright struct {
    23  	PW      *playwright.Playwright
    24  	Browser playwright.Browser
    25  	Page    playwright.Page
    26  }
    27  
    28  func (pw *Playwright) RestartBrowser() (Playwright, error) {
    29  	if err := pw.Browser.Close(); err != nil {
    30  		return Playwright{}, fmt.Errorf("failed to close Playwright browser: %w", err)
    31  	}
    32  	return startPlaywright(pw.PW)
    33  }
    34  
    35  func (pw *Playwright) Cleanup() error {
    36  	if err := pw.Browser.Close(); err != nil {
    37  		return fmt.Errorf("failed to close Playwright browser: %w", err)
    38  	}
    39  	if err := pw.PW.Stop(); err != nil {
    40  		return fmt.Errorf("failed to stop Playwright: %w", err)
    41  	}
    42  	return nil
    43  }
    44  
    45  func startPlaywright(pw *playwright.Playwright) (Playwright, error) {
    46  	browser, err := pw.Chromium.Launch()
    47  	if err != nil {
    48  		return Playwright{}, fmt.Errorf("failed to launch Chromium: %w", err)
    49  	}
    50  	context, err := browser.NewContext()
    51  	if err != nil {
    52  		return Playwright{}, fmt.Errorf("failed to start new Playwright browser context: %w", err)
    53  	}
    54  	page, err := context.NewPage()
    55  	if err != nil {
    56  		return Playwright{}, fmt.Errorf("failed to start new Playwright page: %w", err)
    57  	}
    58  	return Playwright{
    59  		PW:      pw,
    60  		Browser: browser,
    61  		Page:    page,
    62  	}, nil
    63  }
    64  
    65  func InitPlaywright() (Playwright, error) {
    66  	err := playwright.Install(&playwright.RunOptions{
    67  		Verbose:  false,
    68  		Browsers: []string{"chromium"},
    69  	})
    70  	if err != nil {
    71  		return Playwright{}, fmt.Errorf("failed to install Playwright: %w", err)
    72  	}
    73  
    74  	pw, err := playwright.Run()
    75  	if err != nil {
    76  		return Playwright{}, fmt.Errorf("failed to run Playwright: %w", err)
    77  	}
    78  	return startPlaywright(pw)
    79  }
    80  
    81  //go:embed generate_png.js
    82  var genPNGScript string
    83  
    84  const pngPrefix = "data:image/png;base64,"
    85  
    86  // ConvertSVG converts the given SVG into a PNG.
    87  // Note that the resulting PNG has 2x the size (width and height) of the original SVG (see generate_png.js)
    88  func ConvertSVG(page playwright.Page, svg []byte) ([]byte, error) {
    89  	encodedSVG := base64.StdEncoding.EncodeToString(svg)
    90  	pngInterface, err := page.Evaluate(genPNGScript, map[string]interface{}{
    91  		"imgString": "data:image/svg+xml;charset=utf-8;base64," + encodedSVG,
    92  		"scale":     int(SCALE),
    93  	})
    94  	if err != nil {
    95  		return nil, fmt.Errorf("failed to generate png: %w", err)
    96  	}
    97  
    98  	pngString := pngInterface.(string)
    99  	if !strings.HasPrefix(pngString, pngPrefix) {
   100  		if len(pngString) > 50 {
   101  			pngString = pngString[0:50] + "..."
   102  		}
   103  		return nil, fmt.Errorf("invalid PNG: %q", pngString)
   104  	}
   105  	splicedPNGString := pngString[len(pngPrefix):]
   106  	return base64.StdEncoding.DecodeString(splicedPNGString)
   107  }
   108  
   109  func AddExif(png []byte) ([]byte, error) {
   110  	// https://pkg.go.dev/github.com/dsoprea/go-png-image-structure/v2?utm_source=godoc#example-ChunkSlice.SetExif
   111  	im, err := exifcommon.NewIfdMappingWithStandard()
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  
   116  	ti := exif.NewTagIndex()
   117  
   118  	ib := exif.NewIfdBuilder(im, ti, exifcommon.IfdStandardIfdIdentity, exifcommon.TestDefaultByteOrder)
   119  
   120  	err = ib.AddStandardWithName("Make", "D2")
   121  	if err != nil {
   122  		return nil, err
   123  	}
   124  
   125  	err = ib.AddStandardWithName("Model", version.Version)
   126  	if err != nil {
   127  		return nil, err
   128  	}
   129  
   130  	pmp := pngstruct.NewPngMediaParser()
   131  	intfc, err := pmp.ParseBytes(png)
   132  	if err != nil {
   133  		return nil, err
   134  	}
   135  	cs := intfc.(*pngstruct.ChunkSlice)
   136  	err = cs.SetExif(ib)
   137  	if err != nil {
   138  		return nil, err
   139  	}
   140  	b := new(bytes.Buffer)
   141  	err = cs.WriteTo(b)
   142  	if err != nil {
   143  		return nil, err
   144  	}
   145  
   146  	return b.Bytes(), nil
   147  }
   148
View as plain text