✨ initial commit
This commit is contained in:
443
generator.go
Normal file
443
generator.go
Normal file
@@ -0,0 +1,443 @@
|
||||
package freezelib
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/landaiqing/freezelib/font"
|
||||
"github.com/landaiqing/freezelib/svg"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/alecthomas/chroma/v2"
|
||||
formatter "github.com/alecthomas/chroma/v2/formatters/svg"
|
||||
"github.com/alecthomas/chroma/v2/lexers"
|
||||
"github.com/alecthomas/chroma/v2/styles"
|
||||
"github.com/beevik/etree"
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
"github.com/charmbracelet/x/ansi"
|
||||
"github.com/charmbracelet/x/cellbuf"
|
||||
"github.com/kanrichan/resvg-go"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultFontSize = 14.0
|
||||
defaultLineHeight = 1.2
|
||||
)
|
||||
|
||||
// Generator handles the core screenshot generation logic
|
||||
type Generator struct {
|
||||
config *Config
|
||||
}
|
||||
|
||||
// NewGenerator creates a new generator with the given configuration
|
||||
func NewGenerator(config *Config) *Generator {
|
||||
if config == nil {
|
||||
config = DefaultConfig()
|
||||
}
|
||||
return &Generator{config: config}
|
||||
}
|
||||
|
||||
// GenerateFromCode generates an SVG from source code
|
||||
func (g *Generator) GenerateFromCode(code, language string) ([]byte, error) {
|
||||
if err := g.config.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("invalid config: %w", err)
|
||||
}
|
||||
|
||||
// Set language if provided
|
||||
if language != "" {
|
||||
g.config.Language = language
|
||||
}
|
||||
|
||||
// Get lexer for the language
|
||||
var lexer chroma.Lexer
|
||||
if g.config.Language != "" {
|
||||
lexer = lexers.Get(g.config.Language)
|
||||
}
|
||||
if lexer == nil {
|
||||
lexer = lexers.Analyse(code)
|
||||
}
|
||||
if lexer == nil {
|
||||
return nil, errors.New("could not determine language for syntax highlighting")
|
||||
}
|
||||
|
||||
return g.generateSVG(code, lexer, false)
|
||||
}
|
||||
|
||||
// GenerateFromFile generates an SVG from a source code file
|
||||
func (g *Generator) GenerateFromFile(filename string) ([]byte, error) {
|
||||
if err := g.config.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("invalid config: %w", err)
|
||||
}
|
||||
|
||||
// Read file content
|
||||
content, err := os.ReadFile(filename)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read file: %w", err)
|
||||
}
|
||||
|
||||
code := string(content)
|
||||
|
||||
// Get lexer from filename
|
||||
lexer := lexers.Get(filename)
|
||||
if lexer == nil {
|
||||
lexer = lexers.Analyse(code)
|
||||
}
|
||||
if lexer == nil {
|
||||
return nil, errors.New("could not determine language for syntax highlighting")
|
||||
}
|
||||
|
||||
return g.generateSVG(code, lexer, false)
|
||||
}
|
||||
|
||||
// GenerateFromANSI generates an SVG from ANSI terminal output
|
||||
func (g *Generator) GenerateFromANSI(ansiOutput string) ([]byte, error) {
|
||||
if err := g.config.Validate(); err != nil {
|
||||
return nil, fmt.Errorf("invalid config: %w", err)
|
||||
}
|
||||
|
||||
// For ANSI output, we use a text lexer but handle ANSI sequences specially
|
||||
strippedInput := ansi.Strip(ansiOutput)
|
||||
it := chroma.Literator(chroma.Token{Type: chroma.Text, Value: strippedInput})
|
||||
|
||||
return g.generateSVGFromIterator(ansiOutput, it, true)
|
||||
}
|
||||
|
||||
// generateSVG is the core SVG generation function
|
||||
func (g *Generator) generateSVG(input string, lexer chroma.Lexer, isAnsi bool) ([]byte, error) {
|
||||
// Create token iterator
|
||||
var it chroma.Iterator
|
||||
var err error
|
||||
if isAnsi {
|
||||
strippedInput := ansi.Strip(input)
|
||||
it = chroma.Literator(chroma.Token{Type: chroma.Text, Value: strippedInput})
|
||||
} else {
|
||||
it, err = chroma.Coalesce(lexer).Tokenise(nil, input)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not tokenize input: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return g.generateSVGFromIterator(input, it, isAnsi)
|
||||
}
|
||||
|
||||
// generateSVGFromIterator generates SVG from a token iterator
|
||||
func (g *Generator) generateSVGFromIterator(input string, it chroma.Iterator, isAnsi bool) ([]byte, error) {
|
||||
config := g.config
|
||||
|
||||
// Calculate scale factor
|
||||
scale := 1.0
|
||||
autoHeight := config.Height == 0
|
||||
autoWidth := config.Width == 0
|
||||
|
||||
// Expand padding and margin
|
||||
expandedMargin := config.expandMargin(scale)
|
||||
expandedPadding := config.expandPadding(scale)
|
||||
|
||||
// Process input based on line selection
|
||||
processedInput := input
|
||||
if len(config.Lines) == 2 {
|
||||
processedInput = cutLines(input, config.Lines)
|
||||
}
|
||||
|
||||
// Handle text wrapping
|
||||
if config.Wrap > 0 {
|
||||
processedInput = cellbuf.Wrap(processedInput, config.Wrap, "")
|
||||
}
|
||||
|
||||
// Get style
|
||||
style, ok := styles.Registry[strings.ToLower(config.Theme)]
|
||||
if !ok || style == nil {
|
||||
style = styles.Get("github") // fallback to github style
|
||||
}
|
||||
|
||||
// Add background color to style if not present
|
||||
if !style.Has(chroma.Background) {
|
||||
var err error
|
||||
style, err = style.Builder().Add(chroma.Background, "bg:"+config.Background).Build()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not add background: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Get font options
|
||||
fontOptions, err := font.FontOptions(config.Font.Family, config.Font.Size, config.Font.Ligatures, config.Font.File)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid font options: %w", err)
|
||||
}
|
||||
|
||||
// Create SVG formatter
|
||||
f := formatter.New(fontOptions...)
|
||||
|
||||
// Format to SVG
|
||||
buf := &bytes.Buffer{}
|
||||
err = f.Format(buf, style, it)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not format to SVG: %w", err)
|
||||
}
|
||||
|
||||
// Parse SVG document
|
||||
doc := etree.NewDocument()
|
||||
_, err = doc.ReadFrom(buf)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not parse SVG: %w", err)
|
||||
}
|
||||
|
||||
elements := doc.ChildElements()
|
||||
if len(elements) < 1 {
|
||||
return nil, errors.New("invalid SVG output")
|
||||
}
|
||||
|
||||
image := elements[0]
|
||||
|
||||
// Calculate dimensions
|
||||
w, h := svg.GetDimensions(image)
|
||||
imageWidth := float64(w) * scale
|
||||
imageHeight := float64(h) * scale
|
||||
|
||||
// Adjust for font size and line height
|
||||
imageHeight *= config.Font.Size / defaultFontSize
|
||||
imageHeight *= config.LineHeight / defaultLineHeight
|
||||
|
||||
terminalWidth := imageWidth
|
||||
terminalHeight := imageHeight
|
||||
|
||||
hPadding := expandedPadding[left] + expandedPadding[right]
|
||||
hMargin := expandedMargin[left] + expandedMargin[right]
|
||||
vMargin := expandedMargin[top] + expandedMargin[bottom]
|
||||
vPadding := expandedPadding[top] + expandedPadding[bottom]
|
||||
|
||||
// Calculate final dimensions
|
||||
if !autoWidth {
|
||||
imageWidth = config.Width
|
||||
terminalWidth = config.Width - hMargin
|
||||
} else {
|
||||
imageWidth += hMargin + hPadding
|
||||
terminalWidth += hPadding
|
||||
}
|
||||
|
||||
if !autoHeight {
|
||||
imageHeight = config.Height
|
||||
terminalHeight = config.Height - vMargin
|
||||
} else {
|
||||
imageHeight += vMargin + vPadding
|
||||
terminalHeight += vPadding
|
||||
}
|
||||
|
||||
// Get terminal background element
|
||||
terminal := image.SelectElement("rect")
|
||||
if terminal == nil {
|
||||
return nil, errors.New("could not find terminal background element")
|
||||
}
|
||||
|
||||
// Add window controls if enabled
|
||||
if config.Window {
|
||||
windowControls := svg.NewWindowControls(5.5*scale, 19.0*scale, 12.0*scale)
|
||||
svg.Move(windowControls, expandedMargin[left], expandedMargin[top])
|
||||
image.AddChild(windowControls)
|
||||
expandedPadding[top] += 15 * scale
|
||||
}
|
||||
|
||||
// Add corner radius
|
||||
if config.Border.Radius > 0 {
|
||||
svg.AddCornerRadius(terminal, config.Border.Radius*scale)
|
||||
}
|
||||
|
||||
// Add shadow
|
||||
if config.Shadow.Blur > 0 || config.Shadow.X > 0 || config.Shadow.Y > 0 {
|
||||
id := "shadow"
|
||||
svg.AddShadow(image, id, config.Shadow.X*scale, config.Shadow.Y*scale, config.Shadow.Blur*scale)
|
||||
terminal.CreateAttr("filter", fmt.Sprintf("url(#%s)", id))
|
||||
}
|
||||
|
||||
// Process text elements
|
||||
textGroup := image.SelectElement("g")
|
||||
if textGroup != nil {
|
||||
textGroup.CreateAttr("font-size", fmt.Sprintf("%.2fpx", config.Font.Size*scale))
|
||||
textGroup.CreateAttr("clip-path", "url(#terminalMask)")
|
||||
text := textGroup.SelectElements("text")
|
||||
|
||||
offsetLine := 0
|
||||
if len(config.Lines) > 0 {
|
||||
offsetLine = config.Lines[0]
|
||||
}
|
||||
|
||||
lineHeight := config.LineHeight * scale
|
||||
|
||||
for i, line := range text {
|
||||
if isAnsi {
|
||||
line.SetText("")
|
||||
}
|
||||
|
||||
// Add line numbers if enabled
|
||||
if config.ShowLineNumbers {
|
||||
ln := etree.NewElement("tspan")
|
||||
ln.CreateAttr("xml:space", "preserve")
|
||||
ln.CreateAttr("fill", style.Get(chroma.LineNumbers).Colour.String())
|
||||
ln.SetText(fmt.Sprintf("%3d ", i+1+offsetLine))
|
||||
line.InsertChildAt(0, ln)
|
||||
}
|
||||
|
||||
// Position the line
|
||||
x := expandedPadding[left] + expandedMargin[left]
|
||||
y := (float64(i+1))*(config.Font.Size*lineHeight) + expandedPadding[top] + expandedMargin[top]
|
||||
|
||||
svg.Move(line, x, y)
|
||||
|
||||
// Remove lines that are outside the visible area
|
||||
if y > imageHeight-expandedMargin[bottom]-expandedPadding[bottom] {
|
||||
textGroup.RemoveChild(line)
|
||||
}
|
||||
}
|
||||
|
||||
// Process ANSI sequences if needed
|
||||
if isAnsi {
|
||||
processANSI(processedInput, text, textGroup, config, scale)
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate auto width based on content
|
||||
if autoWidth {
|
||||
tabWidth := 4
|
||||
if isAnsi {
|
||||
tabWidth = 6
|
||||
}
|
||||
strippedInput := ansi.Strip(processedInput)
|
||||
longestLine := lipgloss.Width(strings.ReplaceAll(strippedInput, "\t", strings.Repeat(" ", tabWidth)))
|
||||
terminalWidth = float64(longestLine+1) * (config.Font.Size / font.GetFontHeightToWidthRatio())
|
||||
terminalWidth *= scale
|
||||
terminalWidth += hPadding
|
||||
imageWidth = terminalWidth + hMargin
|
||||
}
|
||||
|
||||
// Add border
|
||||
if config.Border.Width > 0 {
|
||||
svg.AddOutline(terminal, config.Border.Width, config.Border.Color)
|
||||
terminalHeight -= config.Border.Width * 2
|
||||
terminalWidth -= config.Border.Width * 2
|
||||
}
|
||||
|
||||
// Adjust for line numbers
|
||||
if config.ShowLineNumbers {
|
||||
if autoWidth {
|
||||
terminalWidth += config.Font.Size * 3 * scale
|
||||
imageWidth += config.Font.Size * 3 * scale
|
||||
} else {
|
||||
terminalWidth -= config.Font.Size * 3
|
||||
}
|
||||
}
|
||||
|
||||
// Add clipping path if needed
|
||||
if !autoHeight || !autoWidth {
|
||||
svg.AddClipPath(image, "terminalMask",
|
||||
expandedMargin[left], expandedMargin[top],
|
||||
terminalWidth, terminalHeight-expandedPadding[bottom])
|
||||
}
|
||||
|
||||
// Set final positions and dimensions
|
||||
svg.Move(terminal, max(expandedMargin[left], config.Border.Width/2), max(expandedMargin[top], config.Border.Width/2))
|
||||
svg.SetDimensions(image, imageWidth, imageHeight)
|
||||
svg.SetDimensions(terminal, terminalWidth, terminalHeight)
|
||||
|
||||
// Convert to bytes
|
||||
return doc.WriteToBytes()
|
||||
}
|
||||
|
||||
// ConvertToPNG converts SVG data to PNG format
|
||||
func (g *Generator) ConvertToPNG(svgData []byte, width, height float64) ([]byte, error) {
|
||||
// Parse SVG document
|
||||
doc := etree.NewDocument()
|
||||
err := doc.ReadFromBytes(svgData)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not parse SVG: %w", err)
|
||||
}
|
||||
|
||||
// Use resvg for conversion
|
||||
worker, err := resvg.NewDefaultWorker(context.Background())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not create resvg worker: %w", err)
|
||||
}
|
||||
defer worker.Close()
|
||||
|
||||
fontdb, err := worker.NewFontDBDefault()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not create font database: %w", err)
|
||||
}
|
||||
defer fontdb.Close()
|
||||
|
||||
// Load embedded fonts
|
||||
if len(font.JetBrainsMonoTTF) > 0 {
|
||||
err = fontdb.LoadFontData(font.JetBrainsMonoTTF)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not load JetBrains Mono font: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
pixmap, err := worker.NewPixmap(uint32(width), uint32(height))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not create pixmap: %w", err)
|
||||
}
|
||||
defer pixmap.Close()
|
||||
|
||||
tree, err := worker.NewTreeFromData(svgData, &resvg.Options{
|
||||
Dpi: 192,
|
||||
ShapeRenderingMode: resvg.ShapeRenderingModeGeometricPrecision,
|
||||
TextRenderingMode: resvg.TextRenderingModeOptimizeLegibility,
|
||||
ImageRenderingMode: resvg.ImageRenderingModeOptimizeQuality,
|
||||
DefaultSizeWidth: float32(width),
|
||||
DefaultSizeHeight: float32(height),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not create SVG tree: %w", err)
|
||||
}
|
||||
defer tree.Close()
|
||||
|
||||
err = tree.ConvertText(fontdb)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not convert text: %w", err)
|
||||
}
|
||||
|
||||
err = tree.Render(resvg.TransformIdentity(), pixmap)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not render SVG: %w", err)
|
||||
}
|
||||
|
||||
pngData, err := pixmap.EncodePNG()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("could not encode PNG: %w", err)
|
||||
}
|
||||
|
||||
return pngData, nil
|
||||
}
|
||||
|
||||
// cutLines cuts the input to the specified line range
|
||||
func cutLines(input string, lines []int) string {
|
||||
if len(lines) != 2 {
|
||||
return input
|
||||
}
|
||||
|
||||
inputLines := strings.Split(input, "\n")
|
||||
start := lines[0]
|
||||
end := lines[1]
|
||||
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
if end >= len(inputLines) || end < 0 {
|
||||
end = len(inputLines) - 1
|
||||
}
|
||||
if start > end {
|
||||
return ""
|
||||
}
|
||||
|
||||
return strings.Join(inputLines[start:end+1], "\n")
|
||||
}
|
||||
|
||||
// max returns the maximum of two float64 values
|
||||
func max(a, b float64) float64 {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
Reference in New Issue
Block a user