chroma-markdown/main.go

package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"

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

// dataPipedIn returns true if the user piped data via stdin.
func dataPipedIn() bool {
	stat, err := os.Stdin.Stat()
	if err != nil {
		return false
	}
	return (stat.Mode() & os.ModeCharDevice) == 0
}

func checkError(err error, msg string) {
	if err != nil {
		if msg != "" {
			fmt.Fprintf(os.Stderr, "Error %s: %v\n", msg, err)
		} else {
			fmt.Fprintf(os.Stderr, "Encountered an error: %v", err)
		}
		os.Exit(2)
	}
}

func init() {
	flag.Usage = func() {
		os.Stderr.WriteString(`usage: chroma-markdown [markdown-file]

`)
		flag.PrintDefaults()
	}
}

func highlight(w io.Writer, source, lexer, style string) error {
	// Determine lexer.
	l := lexers.Get(lexer)
	if l == nil {
		l = lexers.Analyse(source)
	}
	if l == nil {
		l = lexers.Fallback
	}
	l = chroma.Coalesce(l)

	// Determine formatter and style, depending on whether a specific style is
	// requested. If no style is requestd, use classes instead of inline css,
	// leaving the CSS handling to the consumer.
	var f *html.Formatter
	var s *chroma.Style
	if style == "" {
		f = html.New(html.WithClasses(true))
		s = new(chroma.Style)
	} else {
		f = html.New()
		s = styles.Get(style)
	}

	it, err := l.Tokenise(nil, source)
	if err != nil {
		return err
	}

	/* HACK: If background color is white, set to a reasonable, light grey
	 * default. */
	if s.Get(chroma.Background).Background.String() == "#ffffff" {
		senr := s.Get(chroma.Background)
		senr.Background = chroma.ParseColour("#f6f8fa");
		s, err = s.Builder().AddEntry(chroma.Background, senr).Build();
		if err != nil {
			return err
		}
	}

	return f.Format(w, s, it)
}

const Version = "0.1.3"

func main() {
	css := flag.String("css", "", "Path to a CSS import to include at the beginning of the output")
	style := flag.String("style", "", "CSS style to use")
	version := flag.Bool("version", false, "Print the version string")
	v := flag.Bool("v", false, "Print the version string")
	flag.Parse()
	if *version || *v {
		fmt.Printf("chroma-markdown version %s\n", Version)
		os.Exit(2)
	}
	var r io.Reader
	if dataPipedIn() {
		r = os.Stdin
	} else {
		if flag.NArg() != 1 {
			flag.Usage()
		}
		file := flag.Arg(0)
		f, err := os.Open(file)
		checkError(err, "opening file")
		defer f.Close()
		r = bufio.NewReader(f)
	}
	out := new(bytes.Buffer)
	currentCodeBlock := new(bytes.Buffer)
	started := false
	bs := bufio.NewScanner(r)
	lang := ""
	needCSS := false
	cmark, lookErr := exec.LookPath("cmark-gfm")
	for bs.Scan() {
		text := bs.Text()
		trimmed := strings.TrimSpace(text)
		if strings.HasPrefix(trimmed, "```") {
			if started {
				// TODO: compile the code block to markdown
				quickErr := highlight(out, currentCodeBlock.String(), lang, *style)
				checkError(quickErr, "highlighting source code")
				started = false
				currentCodeBlock.Reset()
				lang = ""
				needCSS = true
				continue
			}
			lang = trimmed[3:]
			started = true
			continue
		}
		if started {
			currentCodeBlock.WriteString(text)
			currentCodeBlock.WriteByte('\n')
		} else {
			/* If cmark is not installed then discount will be used, and we have to
			 * do some encoding ourselves. */
			if lookErr != nil {
				text = strings.Replace(text, "\\<", "&lt;", -1)
				text = strings.Replace(text, "\\>", "&gt;", -1)
			}
			out.WriteString(text)
			out.WriteByte('\n')
		}
	}
	checkError(bs.Err(), "reading markdown file")
	f, err := os.CreateTemp("", "chroma-markdown-")
	checkError(err, "creating temporary file")
	w := bufio.NewWriter(f)
	if needCSS && *css != "" {
		_, writeErr := fmt.Fprintf(w, `<link rel="stylesheet" type="text/css" href="%s" />\n`, *css)
		checkError(writeErr, "writing data to temporary file")
	}
	_, writeErr := f.Write(out.Bytes())
	checkError(writeErr, "writing data to temporary file")
	// shell out to markdown because of
	// https://github.com/russross/blackfriday/issues/403
	var args []string
	if lookErr != nil {
		cmark, lookErr = exec.LookPath("markdown")
		checkError(lookErr, "finding markdown binary")
		args = []string{cmark, "-f", "smarty,pants", f.Name()}
	} else {
		args = []string{cmark, "--smart", "--unsafe", "--extension", "table", f.Name()}
	}
	execErr := localExec(cmark, args, []string{})
	checkError(execErr, "executing markdown binary")
	if err := f.Close(); err != nil {
		checkError(err, "closing file")
	}
}