Merge pull request cli#4833 from cli/diff-color

pr diff color output fixes

pkg/cmd/pr/diff/diff.go

@@ -26,7 +26,7 @@ type DiffOptions struct {
 	Finder shared.PRFinder
 
 	SelectorArg string
-	UseColor    string
+	UseColor    bool
 	Patch       bool
 }
 
@@ -36,6 +36,8 @@ func NewCmdDiff(f *cmdutil.Factory, runF func(*DiffOptions) error) *cobra.Comman
 		HttpClient: f.HttpClient,
 	}
 
+	var colorFlag string
+
 	cmd := &cobra.Command{
 		Use:   "diff [<number> | <url> | <branch>]",
 		Short: "View changes in a pull request",
@@ -50,19 +52,22 @@ func NewCmdDiff(f *cmdutil.Factory, runF func(*DiffOptions) error) *cobra.Comman
 			opts.Finder = shared.NewFinder(f)
 
 			if repoOverride, _ := cmd.Flags().GetString("repo"); repoOverride != "" && len(args) == 0 {
-				return cmdutil.FlagErrorf("argument required when using the --repo flag")
+				return cmdutil.FlagErrorf("argument required when using the `--repo` flag")
 			}
 
 			if len(args) > 0 {
 				opts.SelectorArg = args[0]
 			}
 
-			if !validColorFlag(opts.UseColor) {
-				return cmdutil.FlagErrorf("did not understand color: %q. Expected one of always, never, or auto", opts.UseColor)
-			}
-
-			if opts.UseColor == "auto" && !opts.IO.IsStdoutTTY() {
-				opts.UseColor = "never"
+			switch colorFlag {
+			case "always":
+				opts.UseColor = true
+			case "auto":
+				opts.UseColor = opts.IO.ColorEnabled()
+			case "never":
+				opts.UseColor = false
+			default:
+				return cmdutil.FlagErrorf("the value for `--color` must be one of \"auto\", \"always\", or \"never\"")
 			}
 
 			if runF != nil {
@@ -72,7 +77,7 @@ func NewCmdDiff(f *cmdutil.Factory, runF func(*DiffOptions) error) *cobra.Comman
 		},
 	}
 
-	cmd.Flags().StringVar(&opts.UseColor, "color", "auto", "Use color in diff output: {always|never|auto}")
+	cmd.Flags().StringVar(&colorFlag, "color", "auto", "Use color in diff output: {always|never|auto}")
 	cmd.Flags().BoolVar(&opts.Patch, "patch", false, "Display diff in patch format")
 
 	return cmd
@@ -105,34 +110,15 @@ func diffRun(opts *DiffOptions) error {
 	}
 	defer opts.IO.StopPager()
 
-	if opts.UseColor == "never" {
+	if !opts.UseColor {
 		_, err = io.Copy(opts.IO.Out, diff)
 		if errors.Is(err, syscall.EPIPE) {
 			return nil
 		}
 		return err
 	}
 
-	diffLines := bufio.NewScanner(diff)
-	for diffLines.Scan() {
-		diffLine := diffLines.Text()
-		switch {
-		case isHeaderLine(diffLine):
-			fmt.Fprintf(opts.IO.Out, "\x1b[1;38m%s\x1b[m\n", diffLine)
-		case isAdditionLine(diffLine):
-			fmt.Fprintf(opts.IO.Out, "\x1b[32m%s\x1b[m\n", diffLine)
-		case isRemovalLine(diffLine):
-			fmt.Fprintf(opts.IO.Out, "\x1b[31m%s\x1b[m\n", diffLine)
-		default:
-			fmt.Fprintln(opts.IO.Out, diffLine)
-		}
-	}
-
-	if err := diffLines.Err(); err != nil {
-		return fmt.Errorf("error reading pull request diff: %w", err)
-	}
-
-	return nil
+	return colorDiffLines(opts.IO.Out, diff)
 }
 
 func fetchDiff(httpClient *http.Client, baseRepo ghrepo.Interface, prNumber int, asPatch bool) (io.ReadCloser, error) {
@@ -165,9 +151,71 @@ func fetchDiff(httpClient *http.Client, baseRepo ghrepo.Interface, prNumber int,
 	return resp.Body, nil
 }
 
+const lineBufferSize = 4096
+
+var (
+	colorHeader   = []byte("\x1b[1;38m")
+	colorAddition = []byte("\x1b[32m")
+	colorRemoval  = []byte("\x1b[31m")
+	colorReset    = []byte("\x1b[m")
+)
+
+func colorDiffLines(w io.Writer, r io.Reader) error {
+	diffLines := bufio.NewReaderSize(r, lineBufferSize)
+	wasPrefix := false
+	needsReset := false
+
+	for {
+		diffLine, isPrefix, err := diffLines.ReadLine()
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			return fmt.Errorf("error reading pull request diff: %w", err)
+		}
+
+		var color []byte
+		if !wasPrefix {
+			if isHeaderLine(diffLine) {
+				color = colorHeader
+			} else if isAdditionLine(diffLine) {
+				color = colorAddition
+			} else if isRemovalLine(diffLine) {
+				color = colorRemoval
+			}
+		}
+
+		if color != nil {
+			if _, err := w.Write(color); err != nil {
+				return err
+			}
+			needsReset = true
+		}
+
+		if _, err := w.Write(diffLine); err != nil {
+			return err
+		}
+
+		if !isPrefix {
+			if needsReset {
+				if _, err := w.Write(colorReset); err != nil {
+					return err
+				}
+				needsReset = false
+			}
+			if _, err := w.Write([]byte{'\n'}); err != nil {
+				return err
+			}
+		}
+		wasPrefix = isPrefix
+	}
+	return nil
+}
+
 var diffHeaderPrefixes = []string{"+++", "---", "diff", "index"}
 
-func isHeaderLine(dl string) bool {
+func isHeaderLine(l []byte) bool {
+	dl := string(l)
 	for _, p := range diffHeaderPrefixes {
 		if strings.HasPrefix(dl, p) {
 			return true
@@ -176,14 +224,10 @@ func isHeaderLine(dl string) bool {
 	return false
 }
 
-func isAdditionLine(dl string) bool {
-	return strings.HasPrefix(dl, "+")
-}
-
-func isRemovalLine(dl string) bool {
-	return strings.HasPrefix(dl, "-")
+func isAdditionLine(l []byte) bool {
+	return len(l) > 0 && l[0] == '+'
 }
 
-func validColorFlag(c string) bool {
-	return c == "auto" || c == "always" || c == "never"
+func isRemovalLine(l []byte) bool {
+	return len(l) > 0 && l[0] == '-'
 }