enhancement: Add --color flag to cerbos compile

cmd/cerbos/compile/compile.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/alecthomas/kong"
 	"github.com/fatih/color"
+	"github.com/pterm/pterm"
 
 	policyv1 "github.com/cerbos/cerbos/api/genpb/cerbos/policy/v1"
 	internalcompile "github.com/cerbos/cerbos/cmd/cerbos/compile/internal/compilation"
@@ -21,6 +22,7 @@ import (
 	"github.com/cerbos/cerbos/cmd/cerbos/compile/internal/verification"
 	"github.com/cerbos/cerbos/internal/compile"
 	"github.com/cerbos/cerbos/internal/engine"
+	"github.com/cerbos/cerbos/internal/outputcolor"
 	"github.com/cerbos/cerbos/internal/printer"
 	"github.com/cerbos/cerbos/internal/schema"
 	"github.com/cerbos/cerbos/internal/storage/disk"
@@ -44,29 +46,39 @@ cerbos compile --run=Delete /path/to/policy/repo
 cerbos compile --skip-tests /path/to/policy/repo
 `
 
-type Cmd struct {
+type Cmd struct { //nolint:govet // Kong prints fields in order, so we don't want to reorder fields to save bytes.
 	Dir           string               `help:"Policy directory" arg:"" required:"" type:"existingdir"`
-	Output        flagset.OutputFormat `help:"Output format (${enum})" default:"tree" enum:"tree,list,json" short:"o"`
+	IgnoreSchemas bool                 `help:"Ignore schemas during compilation"`
 	Tests         string               `help:"Path to the directory containing tests. Defaults to policy directory." type:"existingdir"`
 	RunRegex      string               `help:"Run only tests that match this regex" name:"run"`
 	SkipTests     bool                 `help:"Skip tests"`
-	IgnoreSchemas bool                 `help:"Ignore schemas during compilation"`
+	Output        flagset.OutputFormat `help:"Output format (${enum})" default:"tree" enum:"tree,list,json" short:"o"`
+	Color         *outputcolor.Level   `help:"Output color level (auto,never,always,256,16m). Defaults to auto." xor:"color"`
+	NoColor       bool                 `help:"Disable colored output" xor:"color"`
 	Verbose       bool                 `help:"Verbose output on test failure"`
-	NoColor       bool                 `help:"Disable colored output"`
 }
 
 func (c *Cmd) Run(k *kong.Kong) error {
 	ctx, stopFunc := signal.NotifyContext(context.Background(), os.Interrupt)
 	defer stopFunc()
 
-	color.NoColor = c.NoColor
+	colorLevel := c.Color.Resolve(c.NoColor)
+
+	color.NoColor = !colorLevel.Enabled()
+
+	if colorLevel.Enabled() {
+		pterm.EnableColor()
+	} else {
+		pterm.DisableColor()
+	}
+
 	p := printer.New(k.Stdout, k.Stderr)
 
 	idx, err := index.Build(ctx, os.DirFS(c.Dir))
 	if err != nil {
 		idxErr := new(index.BuildError)
 		if errors.As(err, &idxErr) {
-			return lint.Display(p, idxErr, c.Output, c.NoColor)
+			return lint.Display(p, idxErr, c.Output, colorLevel)
 		}
 
 		return fmt.Errorf("failed to open directory %s: %w", c.Dir, err)
@@ -83,7 +95,7 @@ func (c *Cmd) Run(k *kong.Kong) error {
 	if err := compile.BatchCompile(idx.GetAllCompilationUnits(ctx), schemaMgr); err != nil {
 		compErr := new(compile.ErrorList)
 		if errors.As(err, compErr) {
-			return internalcompile.Display(p, *compErr, c.Output, c.NoColor)
+			return internalcompile.Display(p, *compErr, c.Output, colorLevel)
 		}
 
 		return fmt.Errorf("failed to create engine: %w", err)
@@ -111,7 +123,7 @@ func (c *Cmd) Run(k *kong.Kong) error {
 			return fmt.Errorf("failed to run tests: %w", err)
 		}
 
-		err = verification.Display(p, results, c.Output, c.Verbose, c.NoColor)
+		err = verification.Display(p, results, c.Output, c.Verbose, colorLevel)
 		if err != nil {
 			return fmt.Errorf("failed to display test results: %w", err)
 		}

cmd/cerbos/compile/internal/compilation/display.go

@@ -7,23 +7,24 @@ import (
 	internalerrors "github.com/cerbos/cerbos/cmd/cerbos/compile/internal/errors"
 	"github.com/cerbos/cerbos/cmd/cerbos/compile/internal/flagset"
 	"github.com/cerbos/cerbos/internal/compile"
+	"github.com/cerbos/cerbos/internal/outputcolor"
 	"github.com/cerbos/cerbos/internal/printer"
 	"github.com/cerbos/cerbos/internal/printer/colored"
 )
 
-func Display(p *printer.Printer, errs compile.ErrorList, output flagset.OutputFormat, noColor bool) error {
+func Display(p *printer.Printer, errs compile.ErrorList, output flagset.OutputFormat, colorLevel outputcolor.Level) error {
 	switch output {
 	case flagset.OutputFormatJSON:
-		return displayJSON(p, errs, noColor)
+		return displayJSON(p, errs, colorLevel)
 	case flagset.OutputFormatList, flagset.OutputFormatTree:
 		return displayList(p, errs)
 	}
 
 	return internalerrors.ErrFailed
 }
 
-func displayJSON(p *printer.Printer, errs compile.ErrorList, noColor bool) error {
-	if err := p.PrintJSON(map[string]compile.ErrorList{"compileErrors": errs}, noColor); err != nil {
+func displayJSON(p *printer.Printer, errs compile.ErrorList, colorLevel outputcolor.Level) error {
+	if err := p.PrintJSON(map[string]compile.ErrorList{"compileErrors": errs}, colorLevel); err != nil {
 		return err
 	}
 

cmd/cerbos/compile/internal/lint/display.go

@@ -6,24 +6,25 @@ package lint
 import (
 	internalerrors "github.com/cerbos/cerbos/cmd/cerbos/compile/internal/errors"
 	"github.com/cerbos/cerbos/cmd/cerbos/compile/internal/flagset"
+	"github.com/cerbos/cerbos/internal/outputcolor"
 	"github.com/cerbos/cerbos/internal/printer"
 	"github.com/cerbos/cerbos/internal/printer/colored"
 	"github.com/cerbos/cerbos/internal/storage/index"
 )
 
-func Display(p *printer.Printer, errs *index.BuildError, output flagset.OutputFormat, noColor bool) error {
+func Display(p *printer.Printer, errs *index.BuildError, output flagset.OutputFormat, colorLevel outputcolor.Level) error {
 	switch output {
 	case flagset.OutputFormatJSON:
-		return displayJSON(p, errs, noColor)
+		return displayJSON(p, errs, colorLevel)
 	case flagset.OutputFormatList, flagset.OutputFormatTree:
 		return displayList(p, errs)
 	}
 
 	return internalerrors.ErrFailed
 }
 
-func displayJSON(p *printer.Printer, errs *index.BuildError, noColor bool) error {
-	if err := p.PrintJSON(map[string]*index.BuildError{"lintErrors": errs}, noColor); err != nil {
+func displayJSON(p *printer.Printer, errs *index.BuildError, colorLevel outputcolor.Level) error {
+	if err := p.PrintJSON(map[string]*index.BuildError{"lintErrors": errs}, colorLevel); err != nil {
 		return err
 	}
 

cmd/cerbos/compile/internal/verification/display.go

@@ -12,14 +12,15 @@ import (
 	policyv1 "github.com/cerbos/cerbos/api/genpb/cerbos/policy/v1"
 	"github.com/cerbos/cerbos/cmd/cerbos/compile/internal/flagset"
 	"github.com/cerbos/cerbos/cmd/cerbos/compile/internal/verification/internal/traces"
+	"github.com/cerbos/cerbos/internal/outputcolor"
 	"github.com/cerbos/cerbos/internal/printer"
 	"github.com/cerbos/cerbos/internal/printer/colored"
 )
 
-func Display(p *printer.Printer, results *policyv1.TestResults, output flagset.OutputFormat, verbose, noColor bool) error {
+func Display(p *printer.Printer, results *policyv1.TestResults, output flagset.OutputFormat, verbose bool, colorLevel outputcolor.Level) error {
 	switch output {
 	case flagset.OutputFormatJSON:
-		return p.PrintProtoJSON(results, noColor)
+		return p.PrintProtoJSON(results, colorLevel)
 	case flagset.OutputFormatTree:
 		return displayTree(p, results, verbose)
 	case flagset.OutputFormatList:

cmd/cerbos/main.go

@@ -10,6 +10,7 @@ import (
 	"github.com/cerbos/cerbos/cmd/cerbos/healthcheck"
 	"github.com/cerbos/cerbos/cmd/cerbos/run"
 	"github.com/cerbos/cerbos/cmd/cerbos/server"
+	"github.com/cerbos/cerbos/internal/outputcolor"
 	"github.com/cerbos/cerbos/internal/util"
 )
 
@@ -27,6 +28,7 @@ func main() {
 		kong.Description("Painless access controls for cloud-native applications"),
 		kong.UsageOnError(),
 		kong.Vars{"version": util.AppVersion()},
+		outputcolor.TypeMapper,
 	)
 
 	ctx.FatalIfErrorf(ctx.Run())

go.mod

@@ -40,6 +40,7 @@ require (
 	github.com/jdxcode/netrc v0.0.0-20210204082910-926c7f70242a
 	github.com/jmoiron/sqlx v1.3.4
 	github.com/jwalton/gchalk v1.2.1
+	github.com/jwalton/go-supportscolor v1.1.0
 	github.com/kavu/go_reuseport v1.5.0
 	github.com/lestrrat-go/jwx v1.2.20
 	github.com/mattn/go-isatty v0.0.14
@@ -156,7 +157,6 @@ require (
 	github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
-	github.com/jwalton/go-supportscolor v1.1.0 // indirect
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
 	github.com/kevinburke/ssh_config v1.1.0 // indirect
 	github.com/klauspost/compress v1.13.6 // indirect

internal/outputcolor/level.go

@@ -0,0 +1,101 @@
+// Copyright 2021-2022 Zenauth Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+package outputcolor
+
+import (
+	"fmt"
+	"os"
+	"reflect"
+
+	"github.com/alecthomas/kong"
+	"github.com/jwalton/go-supportscolor"
+)
+
+type Level uint8
+
+const (
+	None    = Level(supportscolor.None)
+	Basic   = Level(supportscolor.Basic)
+	Ansi256 = Level(supportscolor.Ansi256)
+	Ansi16m = Level(supportscolor.Ansi16m)
+)
+
+var TypeMapper = kong.TypeMapper(reflect.TypeOf((*Level)(nil)), kong.MapperFunc(decode))
+
+func (l *Level) Resolve(disable bool) Level {
+	if disable {
+		return None
+	}
+
+	if l != nil {
+		return *l
+	}
+
+	return Level(supportscolor.SupportsColor(os.Stdout.Fd(), supportscolor.SniffFlagsOption(false)).Level)
+}
+
+func (l Level) Enabled() bool {
+	return l > None
+}
+
+func decode(ctx *kong.DecodeContext, target reflect.Value) error {
+	level, err := scan(ctx)
+	if err != nil {
+		return err
+	}
+
+	target.Set(reflect.ValueOf(level))
+	return nil
+}
+
+func scan(ctx *kong.DecodeContext) (*Level, error) {
+	token := ctx.Scan.Peek()
+
+	switch token.Type {
+	case kong.FlagValueToken:
+		return parse(ctx.Scan.Pop().Value)
+
+	case kong.ShortFlagTailToken, kong.UntypedToken:
+		level, err := parse(token.Value)
+		if err == nil {
+			ctx.Scan.Pop()
+			return level, nil
+		}
+
+	default:
+	}
+
+	return pointer(Basic), nil
+}
+
+func parse(v interface{}) (*Level, error) {
+	s, ok := v.(string)
+	if !ok {
+		return nil, fmt.Errorf("invalid flag value (expected string, got %T)", v)
+	}
+
+	switch s {
+	case "auto":
+		return nil, nil
+
+	case "false", "never":
+		return pointer(None), nil
+
+	case "true", "always":
+		return pointer(Basic), nil
+
+	case "256":
+		return pointer(Ansi256), nil
+
+	case "16m", "full", "truecolor":
+		return pointer(Ansi16m), nil
+
+	default:
+		return nil, fmt.Errorf("invalid value for output color level: %q", s)
+	}
+}
+
+func pointer(level Level) *Level {
+	return &level
+}