cmd: --fail earlier without reading entire input files (#502)

Author: scop

Diff for: cmd/chroma/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -82,6 +83,54 @@ type nopFlushableWriter struct{ io.Writer }
 
 func (n *nopFlushableWriter) Flush() error { return nil }
 
+// prepareLenient prepares contents and lexer for input, using fallback lexer if no specific one is available for it.
+func prepareLenient(ctx *kong.Context, r io.Reader, filename string) (string, chroma.Lexer) {
+	data, err := ioutil.ReadAll(r)
+	ctx.FatalIfErrorf(err)
+
+	contents := string(data)
+	lexer := selexer(filename, contents)
+	if lexer == nil {
+		lexer = lexers.Fallback
+	}
+
+	return contents, lexer
+}
+
+// prepareSpecific prepares contents and lexer for input, exiting if there is no specific lexer available for it.
+// Input is consumed only up to peekSize for lexer selection. With fullSize -1, consume r until EOF.
+func prepareSpecific(ctx *kong.Context, r io.Reader, filename string, peekSize, fullSize int) (string, chroma.Lexer) {
+	data := make([]byte, peekSize)
+	n, err := io.ReadFull(r, data)
+	if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) {
+		ctx.FatalIfErrorf(err)
+	}
+
+	lexer := selexer(filename, string(data[:n]))
+	if lexer == nil {
+		ctx.Exit(1)
+	}
+
+	if n < peekSize {
+		return string(data[:n]), lexer
+	}
+	var ndata []byte
+	if fullSize == -1 {
+		rest, err := io.ReadAll(r)
+		ctx.FatalIfErrorf(err)
+		ndata = make([]byte, n + len(rest))
+		copy(ndata, data[:n])
+		copy(ndata[n:], rest)
+	} else {
+		ndata = make([]byte, fullSize)
+		copy(ndata, data[:n])
+		_, err = io.ReadFull(r, ndata[n:])
+		ctx.FatalIfErrorf(err)
+	}
+
+	return string(ndata), lexer
+}
+
 func main() {
 	ctx := kong.Parse(&cli, kong.Description(description), kong.Vars{
 		"version":    fmt.Sprintf("%s-%s-%s", version, commit, date),
@@ -160,18 +209,37 @@ func main() {
 		configureHTMLFormatter(ctx)
 	}
 	if len(cli.Files) == 0 {
-		contents, err := ioutil.ReadAll(os.Stdin)
-		ctx.FatalIfErrorf(err)
-		format(ctx, w, style, lex(ctx, cli.Filename, string(contents)))
+		var contents string
+		var lexer chroma.Lexer
+		if cli.Fail {
+			contents, lexer = prepareSpecific(ctx, os.Stdin, cli.Filename, 1024, -1)
+		} else {
+			contents, lexer = prepareLenient(ctx, os.Stdin, cli.Filename)
+		}
+		format(ctx, w, style, lex(ctx, lexer, contents))
 	} else {
 		for _, filename := range cli.Files {
-			contents, err := ioutil.ReadFile(filename)
+			file, err := os.Open(filename)
 			ctx.FatalIfErrorf(err)
+
 			if cli.Check {
-				check(filename, lex(ctx, filename, string(contents)))
+				contents, lexer := prepareLenient(ctx, file, filename)
+				check(filename, lex(ctx, lexer, contents))
 			} else {
-				format(ctx, w, style, lex(ctx, filename, string(contents)))
+				var contents string
+				var lexer chroma.Lexer
+				if cli.Fail {
+					fi, err := file.Stat()
+					ctx.FatalIfErrorf(err)
+					contents, lexer = prepareSpecific(ctx, file, filename, 1024, int(fi.Size()))
+				} else {
+					contents, lexer = prepareLenient(ctx, file, filename)
+				}
+				format(ctx, w, style, lex(ctx, lexer, contents))
 			}
+
+			err = file.Close()
+			ctx.FatalIfErrorf(err)
 		}
 	}
 }
@@ -241,14 +309,7 @@ func listAll() {
 	fmt.Println()
 }
 
-func lex(ctx *kong.Context, path string, contents string) chroma.Iterator {
-	lexer := selexer(path, contents)
-	if lexer == nil {
-		if cli.Fail {
-			ctx.Exit(1)
-		}
-		lexer = lexers.Fallback
-	}
+func lex(ctx *kong.Context, lexer chroma.Lexer, contents string) chroma.Iterator {
 	if rel, ok := lexer.(*chroma.RegexLexer); ok {
 		rel.Trace(cli.Trace)
 	}