kitty/tools/cmd/diff/highlight.go

// License: GPLv3 Copyright: 2023, Kovid Goyal, <kovid at kovidgoyal.net>

package diff

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"kitty/tools/utils"
	"kitty/tools/utils/images"

	"github.com/alecthomas/chroma/v2"
	"github.com/alecthomas/chroma/v2/lexers"
	"github.com/alecthomas/chroma/v2/styles"
)

var _ = fmt.Print
var _ = os.WriteFile

var ErrNoLexer = errors.New("No lexer available for this format")
var DefaultStyle = (&utils.Once[*chroma.Style]{Run: func() *chroma.Style {
	// Default style generated by python style.py default pygments.styles.default.DefaultStyle
	// with https://raw.githubusercontent.com/alecthomas/chroma/master/_tools/style.py
	return styles.Register(chroma.MustNewStyle("default", chroma.StyleEntries{
		chroma.TextWhitespace:        "#bbbbbb",
		chroma.Comment:               "italic #3D7B7B",
		chroma.CommentPreproc:        "noitalic #9C6500",
		chroma.Keyword:               "bold #008000",
		chroma.KeywordPseudo:         "nobold",
		chroma.KeywordType:           "nobold #B00040",
		chroma.Operator:              "#666666",
		chroma.OperatorWord:          "bold #AA22FF",
		chroma.NameBuiltin:           "#008000",
		chroma.NameFunction:          "#0000FF",
		chroma.NameClass:             "bold #0000FF",
		chroma.NameNamespace:         "bold #0000FF",
		chroma.NameException:         "bold #CB3F38",
		chroma.NameVariable:          "#19177C",
		chroma.NameConstant:          "#880000",
		chroma.NameLabel:             "#767600",
		chroma.NameEntity:            "bold #717171",
		chroma.NameAttribute:         "#687822",
		chroma.NameTag:               "bold #008000",
		chroma.NameDecorator:         "#AA22FF",
		chroma.LiteralString:         "#BA2121",
		chroma.LiteralStringDoc:      "italic",
		chroma.LiteralStringInterpol: "bold #A45A77",
		chroma.LiteralStringEscape:   "bold #AA5D1F",
		chroma.LiteralStringRegex:    "#A45A77",
		chroma.LiteralStringSymbol:   "#19177C",
		chroma.LiteralStringOther:    "#008000",
		chroma.LiteralNumber:         "#666666",
		chroma.GenericHeading:        "bold #000080",
		chroma.GenericSubheading:     "bold #800080",
		chroma.GenericDeleted:        "#A00000",
		chroma.GenericInserted:       "#008400",
		chroma.GenericError:          "#E40000",
		chroma.GenericEmph:           "italic",
		chroma.GenericStrong:         "bold",
		chroma.GenericPrompt:         "bold #000080",
		chroma.GenericOutput:         "#717171",
		chroma.GenericTraceback:      "#04D",
		chroma.Error:                 "border:#FF0000",
		chroma.Background:            " bg:#f8f8f8",
	}))
}}).Get

// Clear the background colour.
func clear_background(style *chroma.Style) *chroma.Style {
	builder := style.Builder()
	bg := builder.Get(chroma.Background)
	bg.Background = 0
	bg.NoInherit = true
	builder.AddEntry(chroma.Background, bg)
	style, _ = builder.Build()
	return style
}

const SGR_PREFIX = "\033["
const SGR_SUFFIX = "m"

func ansi_formatter(w io.Writer, style *chroma.Style, it chroma.Iterator) error {
	style = clear_background(style)
	before, after := make([]byte, 0, 64), make([]byte, 0, 64)
	nl := []byte{'\n'}
	write_sgr := func(which []byte) {
		if len(which) > 1 {
			w.Write(utils.UnsafeStringToBytes(SGR_PREFIX))
			w.Write(which[:len(which)-1])
			w.Write(utils.UnsafeStringToBytes(SGR_SUFFIX))
		}
	}
	write := func(text string) {
		write_sgr(before)
		w.Write(utils.UnsafeStringToBytes(text))
		write_sgr(after)
	}

	for token := it(); token != chroma.EOF; token = it() {
		entry := style.Get(token.Type)
		before, after = before[:0], after[:0]
		if !entry.IsZero() {
			if entry.Bold == chroma.Yes {
				before = append(before, '1', ';')
				after = append(after, '2', '2', '1', ';')
			}
			if entry.Underline == chroma.Yes {
				before = append(before, '4', ';')
				after = append(after, '2', '4', ';')
			}
			if entry.Italic == chroma.Yes {
				before = append(before, '3', ';')
				after = append(after, '2', '3', ';')
			}
			if entry.Colour.IsSet() {
				before = append(before, fmt.Sprintf("38:2:%d:%d:%d;", entry.Colour.Red(), entry.Colour.Green(), entry.Colour.Blue())...)
				after = append(after, '3', '9', ';')
			}
		}
		// independently format each line in a multiline token, needed for the diff kitten highlighting to work, also
		// pagers like less reset SGR formatting at line boundaries
		text := token.Value
		for text != "" {
			idx := strings.IndexByte(text, '\n')
			if idx < 0 {
				write(text)
				break
			}
			write(text[:idx])
			w.Write(nl)
			text = text[idx+1:]
		}
	}
	return nil
}

func highlight_file(path string) (highlighted string, err error) {
	filename_for_detection := filepath.Base(path)
	ext := filepath.Ext(filename_for_detection)
	if ext != "" {
		ext = strings.ToLower(ext[1:])
		r := conf.Syntax_aliases[ext]
		if r != "" {
			filename_for_detection = "file." + r
		}
	}
	text, err := data_for_path(path)
	if err != nil {
		return "", err
	}
	text = sanitize_control_codes(text)
	lexer := lexers.Match(filename_for_detection)
	if lexer == nil {
		if err == nil {
			lexer = lexers.Analyse(text)
		}
	}
	if lexer == nil {
		return "", fmt.Errorf("Cannot highlight %#v: %w", path, ErrNoLexer)
	}
	lexer = chroma.Coalesce(lexer)
	name := conf.Pygments_style
	var style *chroma.Style
	if name == "default" {
		style = DefaultStyle()
	} else {
		style = styles.Get(name)
	}
	if style == nil {
		if conf.Background.IsDark() && !conf.Foreground.IsDark() {
			style = styles.Get("monokai")
			if style == nil {
				style = styles.Get("github-dark")
			}
		} else {
			style = DefaultStyle()
		}
		if style == nil {
			style = styles.Fallback
		}
	}
	iterator, err := lexer.Tokenise(nil, text)
	if err != nil {
		return "", err
	}
	formatter := chroma.FormatterFunc(ansi_formatter)
	w := strings.Builder{}
	w.Grow(len(text) * 2)
	err = formatter.Format(&w, style, iterator)
	// os.WriteFile(path+".highlighted", []byte(w.String()), 0o600)
	return w.String(), err
}

func highlight_all(paths []string) {
	ctx := images.Context{}
	ctx.Parallel(0, len(paths), func(nums <-chan int) {
		for i := range nums {
			path := paths[i]
			raw, err := highlight_file(path)
			if err == nil {
				highlighted_lines_cache.Set(path, text_to_lines(sanitize_tabs_and_carriage_returns(raw)))
			}
		}
	})
}