Skip to content

Commit

Permalink
enhancement: Add --color flag to cerbos compile
Browse files Browse the repository at this point in the history
By default, the output color level is detected from the environment by checking
if stdout is a TTY and looking for various environment variables.

See https://pkg.go.dev/github.com/jwalton/go-supportscolor for details.

Signed-off-by: Andrew Haines <haines@cerbos.dev>
  • Loading branch information
haines committed Mar 18, 2022
1 parent c67970d commit 058cf66
Show file tree
Hide file tree
Showing 8 changed files with 146 additions and 34 deletions.
28 changes: 20 additions & 8 deletions cmd/cerbos/compile/compile.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand All @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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)
}
Expand Down
9 changes: 5 additions & 4 deletions cmd/cerbos/compile/internal/compilation/display.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
9 changes: 5 additions & 4 deletions cmd/cerbos/compile/internal/lint/display.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
5 changes: 3 additions & 2 deletions cmd/cerbos/compile/internal/verification/display.go
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 2 additions & 0 deletions cmd/cerbos/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand All @@ -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())
Expand Down
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
93 changes: 93 additions & 0 deletions internal/outputcolor/level.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
// 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 {
value, err := scan(ctx)
if err != nil {
return err
}

level, err := parse(value)
if err != nil {
return err
}

target.Set(reflect.ValueOf(level))
return nil
}

func scan(ctx *kong.DecodeContext) (string, error) {
if ctx.Scan.Peek().Type != kong.FlagValueToken {
return "true", nil
}

value := ctx.Scan.Pop().Value
if result, ok := value.(string); ok {
return result, nil
}

return "", fmt.Errorf("invalid flag value (expected string, got %T)", value)
}

func parse(value string) (*Level, error) {
var level Level
switch value {
case "auto":
return nil, nil

case "false", "never":
level = None

case "true", "always":
level = Basic

case "256":
level = Ansi256

case "16m", "full", "truecolor":
level = Ansi16m

default:
return nil, fmt.Errorf("invalid value for output color level: %q", value)
}

return &level, nil
}
32 changes: 17 additions & 15 deletions internal/printer/printer.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ import (
"github.com/alecthomas/chroma/styles"
effectv1 "github.com/cerbos/cerbos/api/genpb/cerbos/effect/v1"
enginev1 "github.com/cerbos/cerbos/api/genpb/cerbos/engine/v1"
"github.com/cerbos/cerbos/internal/outputcolor"
"github.com/cerbos/cerbos/internal/printer/colored"
"github.com/jwalton/gchalk"
"google.golang.org/protobuf/encoding/protojson"
"google.golang.org/protobuf/proto"
)
Expand All @@ -38,20 +38,22 @@ func (p *Printer) Printf(format string, args ...interface{}) {
fmt.Fprintf(p.stdout, format, args...)
}

func (p *Printer) coloredJSON(data string) error {
func (p *Printer) coloredJSON(data string, colorLevel outputcolor.Level) error {
lexer := chroma.Coalesce(lexers.Get("json"))
if lexer == nil {
lexer = lexers.Fallback
}

var formatter chroma.Formatter
switch gchalk.GetLevel() {
case gchalk.LevelAnsi256:
switch colorLevel {
case outputcolor.Basic:
formatter = formatters.TTY
case outputcolor.Ansi256:
formatter = formatters.TTY256
case gchalk.LevelAnsi16m:
case outputcolor.Ansi16m:
formatter = formatters.TTY16m
default:
formatter = formatters.TTY
formatter = formatters.NoOp
}

iterator, err := lexer.Tokenise(nil, data)
Expand All @@ -62,37 +64,37 @@ func (p *Printer) coloredJSON(data string) error {
return formatter.Format(p.stdout, styles.SolarizedDark256, iterator)
}

func (p *Printer) PrintJSON(val interface{}, noColor bool) error {
func (p *Printer) PrintJSON(val interface{}, colorLevel outputcolor.Level) error {
var data bytes.Buffer
var enc *json.Encoder
if noColor {
enc = json.NewEncoder(p.stdout)
} else {
if colorLevel.Enabled() {
enc = json.NewEncoder(&data)
} else {
enc = json.NewEncoder(p.stdout)
}

enc.SetIndent("", " ")
if err := enc.Encode(val); err != nil {
return fmt.Errorf("failed to encode JSON: %w", err)
}

if !noColor {
return p.coloredJSON(data.String())
if colorLevel.Enabled() {
return p.coloredJSON(data.String(), colorLevel)
}

return nil
}

func (p *Printer) PrintProtoJSON(message proto.Message, noColor bool) error {
func (p *Printer) PrintProtoJSON(message proto.Message, colorLevel outputcolor.Level) error {
data, err := protojson.MarshalOptions{Multiline: true}.Marshal(message)
if err != nil {
return fmt.Errorf("failed to encode JSON: %w", err)
}

output := fmt.Sprintf("%s\n", data)

if !noColor {
return p.coloredJSON(output)
if colorLevel.Enabled() {
return p.coloredJSON(output, colorLevel)
}

fmt.Fprint(p.stdout, output)
Expand Down

0 comments on commit 058cf66

Please sign in to comment.