Merge pull request #1181 from hashicorp/f-plan-cli

cli: nomad plan command

command/meta.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/hashicorp/nomad/api"
 	"github.com/mitchellh/cli"
+	"github.com/mitchellh/colorstring"
 )
 
 const (
@@ -38,6 +39,9 @@ type Meta struct {
 
 	// These are set by the command line flags.
 	flagAddress string
+
+	// Whether to not-colorize output
+	noColor bool
 }
 
 // FlagSet returns a FlagSet with the common flags that every
@@ -51,6 +55,7 @@ func (m *Meta) FlagSet(n string, fs FlagSetFlags) *flag.FlagSet {
 	// client connectivity options.
 	if fs&FlagSetClient != 0 {
 		f.StringVar(&m.flagAddress, "address", "", "")
+		f.BoolVar(&m.noColor, "no-color", false, "")
 	}
 
 	// Create an io.Writer that writes to our UI properly for errors.
@@ -82,6 +87,14 @@ func (m *Meta) Client() (*api.Client, error) {
 	return api.NewClient(config)
 }
 
+func (m *Meta) Colorize() *colorstring.Colorize {
+	return &colorstring.Colorize{
+		Colors:  colorstring.DefaultColors,
+		Disable: m.noColor,
+		Reset:   true,
+	}
+}
+
 // generalOptionsUsage returns the help string for the global options.
 func generalOptionsUsage() string {
 	helpText := `

command/meta_test.go

@@ -18,7 +18,7 @@ func TestMeta_FlagSet(t *testing.T) {
 		},
 		{
 			FlagSetClient,
-			[]string{"address"},
+			[]string{"address", "no-color"},
 		},
 	}
 

command/plan.go

@@ -0,0 +1,441 @@
+package command
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/hashicorp/nomad/api"
+	"github.com/hashicorp/nomad/jobspec"
+	"github.com/hashicorp/nomad/scheduler"
+	"github.com/mitchellh/colorstring"
+)
+
+const (
+	jobModifyIndexHelp = `To submit the job with version verification run:
+
+nomad run -check-index %d %s
+
+When running the job with the check-index flag, the job will only be run if the
+server side version matches the the job modify index returned. If the index has
+changed, another user has modified the job and the plan's results are
+potentially invalid.`
+)
+
+type PlanCommand struct {
+	Meta
+	color *colorstring.Colorize
+}
+
+func (c *PlanCommand) Help() string {
+	helpText := `
+Usage: nomad plan [options] <file>
+
+  Plan invokes a dry-run of the scheduler to determine the effects of submitting
+  either a new or updated version of a job. The plan will not result in any
+  changes to the cluster but gives insight into whether the job could be run
+  successfully and how it would affect existing allocations.
+
+  A job modify index is returned with the plan. This value can be used when
+  submitting the job using "nomad run -check-index", which will check that the job
+  was not modified between the plan and run command before invoking the
+  scheduler. This ensures the job has not been modified since the plan.
+
+  A structured diff between the local and remote job is displayed to
+  give insight into what the scheduler will attempt to do and why.
+
+General Options:
+
+  ` + generalOptionsUsage() + `
+
+Run Options:
+
+  -diff
+    Defaults to true, but can be toggled off to omit diff output.
+
+  -no-color
+    Disable colored output.
+
+  -verbose
+    Increase diff verbosity.
+`
+	return strings.TrimSpace(helpText)
+}
+
+func (c *PlanCommand) Synopsis() string {
+	return "Dry-run a job update to determine its effects"
+}
+
+func (c *PlanCommand) Run(args []string) int {
+	var diff, verbose bool
+
+	flags := c.Meta.FlagSet("plan", FlagSetClient)
+	flags.Usage = func() { c.Ui.Output(c.Help()) }
+	flags.BoolVar(&diff, "diff", true, "")
+	flags.BoolVar(&verbose, "verbose", false, "")
+
+	if err := flags.Parse(args); err != nil {
+		return 1
+	}
+
+	// Check that we got exactly one job
+	args = flags.Args()
+	if len(args) != 1 {
+		c.Ui.Error(c.Help())
+		return 1
+	}
+	file := args[0]
+
+	// Parse the job file
+	job, err := jobspec.ParseFile(file)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error parsing job file %s: %s", file, err))
+		return 1
+	}
+
+	// Initialize any fields that need to be.
+	job.InitFields()
+
+	// Check that the job is valid
+	if err := job.Validate(); err != nil {
+		c.Ui.Error(fmt.Sprintf("Error validating job: %s", err))
+		return 1
+	}
+
+	// Convert it to something we can use
+	apiJob, err := convertStructJob(job)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error converting job: %s", err))
+		return 1
+	}
+
+	// Get the HTTP client
+	client, err := c.Meta.Client()
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
+		return 1
+	}
+
+	// Submit the job
+	resp, _, err := client.Jobs().Plan(apiJob, diff, nil)
+	if err != nil {
+		c.Ui.Error(fmt.Sprintf("Error during plan: %s", err))
+		return 1
+	}
+
+	// Print the diff if not disabled
+	if diff {
+		c.Ui.Output(fmt.Sprintf("%s\n",
+			c.Colorize().Color(strings.TrimSpace(formatJobDiff(resp.Diff, verbose)))))
+	}
+
+	// Print the scheduler dry-run output
+	c.Ui.Output(c.Colorize().Color("[bold]Scheduler dry-run:[reset]"))
+	c.Ui.Output(c.Colorize().Color(formatDryRun(resp.CreatedEvals)))
+
+	// Print the job index info
+	c.Ui.Output(c.Colorize().Color(formatJobModifyIndex(resp.JobModifyIndex, file)))
+	return 0
+}
+
+// formatJobModifyIndex produces a help string that displays the job modify
+// index and how to submit a job with it.
+func formatJobModifyIndex(jobModifyIndex uint64, jobName string) string {
+	help := fmt.Sprintf(jobModifyIndexHelp, jobModifyIndex, jobName)
+	out := fmt.Sprintf("[reset][bold]Job Modify Index: %d[reset]\n%s", jobModifyIndex, help)
+	return out
+}
+
+// formatDryRun produces a string explaining the results of the dry run.
+func formatDryRun(evals []*api.Evaluation) string {
+	var rolling *api.Evaluation
+	var blocked *api.Evaluation
+	for _, eval := range evals {
+		if eval.TriggeredBy == "rolling-update" {
+			rolling = eval
+		} else if eval.Status == "blocked" {
+			blocked = eval
+		}
+	}
+
+	var out string
+	if blocked == nil {
+		out = "[bold][green]  - All tasks successfully allocated.[reset]\n"
+	} else {
+		out = "[bold][yellow]  - WARNING: Failed to place all allocations.[reset]\n"
+	}
+
+	if rolling != nil {
+		out += fmt.Sprintf("[green]  - Rolling update, next evaluation will be in %s.\n", rolling.Wait)
+	}
+
+	return out
+}
+
+// formatJobDiff produces an annoted diff of the the job. If verbose mode is
+// set, added or deleted task groups and tasks are expanded.
+func formatJobDiff(job *api.JobDiff, verbose bool) string {
+	marker, _ := getDiffString(job.Type)
+	out := fmt.Sprintf("%s[bold]Job: %q\n", marker, job.ID)
+
+	// Determine the longest markers and fields so that the output can be
+	// properly alligned.
+	longestField, longestMarker := getLongestPrefixes(job.Fields, job.Objects)
+	for _, tg := range job.TaskGroups {
+		if _, l := getDiffString(tg.Type); l > longestMarker {
+			longestMarker = l
+		}
+	}
+
+	// Only show the job's field and object diffs if the job is edited or
+	// verbose mode is set.
+	if job.Type == "Edited" || verbose {
+		fo := alignedFieldAndObjects(job.Fields, job.Objects, 0, longestField, longestMarker)
+		out += fo
+		if len(fo) > 0 {
+			out += "\n"
+		}
+	}
+
+	// Print the task groups
+	for _, tg := range job.TaskGroups {
+		_, mLength := getDiffString(tg.Type)
+		kPrefix := longestMarker - mLength
+		out += fmt.Sprintf("%s\n", formatTaskGroupDiff(tg, kPrefix, verbose))
+	}
+
+	return out
+}
+
+// formatTaskGroupDiff produces an annotated diff of a task group. If the
+// verbose field is set, the task groups fields and objects are expanded even if
+// the full object is an addition or removal. tgPrefix is the number of spaces to prefix
+// the output of the task group.
+func formatTaskGroupDiff(tg *api.TaskGroupDiff, tgPrefix int, verbose bool) string {
+	marker, _ := getDiffString(tg.Type)
+	out := fmt.Sprintf("%s%s[bold]Task Group: %q[reset]", marker, strings.Repeat(" ", tgPrefix), tg.Name)
+
+	// Append the updates and colorize them
+	if l := len(tg.Updates); l > 0 {
+		updates := make([]string, 0, l)
+		for updateType, count := range tg.Updates {
+			var color string
+			switch updateType {
+			case scheduler.UpdateTypeIgnore:
+			case scheduler.UpdateTypeCreate:
+				color = "[green]"
+			case scheduler.UpdateTypeDestroy:
+				color = "[red]"
+			case scheduler.UpdateTypeMigrate:
+				color = "[blue]"
+			case scheduler.UpdateTypeInplaceUpdate:
+				color = "[cyan]"
+			case scheduler.UpdateTypeDestructiveUpdate:
+				color = "[yellow]"
+			}
+			updates = append(updates, fmt.Sprintf("[reset]%s%d %s", color, count, updateType))
+		}
+		out += fmt.Sprintf(" (%s[reset])\n", strings.Join(updates, ", "))
+	} else {
+		out += "[reset]\n"
+	}
+
+	// Determine the longest field and markers so the output is properly
+	// alligned
+	longestField, longestMarker := getLongestPrefixes(tg.Fields, tg.Objects)
+	for _, task := range tg.Tasks {
+		if _, l := getDiffString(task.Type); l > longestMarker {
+			longestMarker = l
+		}
+	}
+
+	// Only show the task groups's field and object diffs if the group is edited or
+	// verbose mode is set.
+	subStartPrefix := tgPrefix + 2
+	if tg.Type == "Edited" || verbose {
+		fo := alignedFieldAndObjects(tg.Fields, tg.Objects, subStartPrefix, longestField, longestMarker)
+		out += fo
+		if len(fo) > 0 {
+			out += "\n"
+		}
+	}
+
+	// Output the tasks
+	for _, task := range tg.Tasks {
+		_, mLength := getDiffString(task.Type)
+		prefix := longestMarker - mLength
+		out += fmt.Sprintf("%s\n", formatTaskDiff(task, subStartPrefix, prefix, verbose))
+	}
+
+	return out
+}
+
+// formatTaskDiff produces an annotated diff of a task. If the verbose field is
+// set, the tasks fields and objects are expanded even if the full object is an
+// addition or removal. startPrefix is the number of spaces to prefix the output of
+// the task and taskPrefix is the number of spaces to put betwen the marker and
+// task name output.
+func formatTaskDiff(task *api.TaskDiff, startPrefix, taskPrefix int, verbose bool) string {
+	marker, _ := getDiffString(task.Type)
+	out := fmt.Sprintf("%s%s%s[bold]Task: %q",
+		strings.Repeat(" ", startPrefix), marker, strings.Repeat(" ", taskPrefix), task.Name)
+	if len(task.Annotations) != 0 {
+		out += fmt.Sprintf(" [reset](%s)", colorAnnotations(task.Annotations))
+	}
+
+	if task.Type == "None" {
+		return out
+	} else if (task.Type == "Deleted" || task.Type == "Added") && !verbose {
+		// Exit early if the job was not edited and it isn't verbose output
+		return out
+	} else {
+		out += "\n"
+	}
+
+	subStartPrefix := startPrefix + 2
+	longestField, longestMarker := getLongestPrefixes(task.Fields, task.Objects)
+	out += alignedFieldAndObjects(task.Fields, task.Objects, subStartPrefix, longestField, longestMarker)
+	return out
+}
+
+// formatObjectDiff produces an annotated diff of an object. startPrefix is the
+// number of spaces to prefix the output of the object and keyPrefix is the number
+// of spaces to put betwen the marker and object name output.
+func formatObjectDiff(diff *api.ObjectDiff, startPrefix, keyPrefix int) string {
+	start := strings.Repeat(" ", startPrefix)
+	marker, _ := getDiffString(diff.Type)
+	out := fmt.Sprintf("%s%s%s%s {\n", start, marker, strings.Repeat(" ", keyPrefix), diff.Name)
+
+	// Determine the length of the longest name and longest diff marker to
+	// properly align names and values
+	longestField, longestMarker := getLongestPrefixes(diff.Fields, diff.Objects)
+	subStartPrefix := startPrefix + 2
+	out += alignedFieldAndObjects(diff.Fields, diff.Objects, subStartPrefix, longestField, longestMarker)
+	return fmt.Sprintf("%s\n%s}", out, start)
+}
+
+// formatFieldDiff produces an annotated diff of a field. startPrefix is the
+// number of spaces to prefix the output of the field, keyPrefix is the number
+// of spaces to put betwen the marker and field name output and valuePrefix is
+// the number of spaces to put infront of the value for aligning values.
+func formatFieldDiff(diff *api.FieldDiff, startPrefix, keyPrefix, valuePrefix int) string {
+	marker, _ := getDiffString(diff.Type)
+	out := fmt.Sprintf("%s%s%s%s: %s",
+		strings.Repeat(" ", startPrefix),
+		marker, strings.Repeat(" ", keyPrefix),
+		diff.Name,
+		strings.Repeat(" ", valuePrefix))
+
+	switch diff.Type {
+	case "Added":
+		out += fmt.Sprintf("%q", diff.New)
+	case "Deleted":
+		out += fmt.Sprintf("%q", diff.Old)
+	case "Edited":
+		out += fmt.Sprintf("%q => %q", diff.Old, diff.New)
+	default:
+		out += fmt.Sprintf("%q", diff.New)
+	}
+
+	// Color the annotations where possible
+	if l := len(diff.Annotations); l != 0 {
+		out += fmt.Sprintf(" (%s)", colorAnnotations(diff.Annotations))
+	}
+
+	return out
+}
+
+// alignedFieldAndObjects is a helper method that prints fields and objects
+// properly aligned.
+func alignedFieldAndObjects(fields []*api.FieldDiff, objects []*api.ObjectDiff,
+	startPrefix, longestField, longestMarker int) string {
+
+	var out string
+	numFields := len(fields)
+	numObjects := len(objects)
+	haveObjects := numObjects != 0
+	for i, field := range fields {
+		_, mLength := getDiffString(field.Type)
+		kPrefix := longestMarker - mLength
+		vPrefix := longestField - len(field.Name)
+		out += formatFieldDiff(field, startPrefix, kPrefix, vPrefix)
+
+		// Avoid a dangling new line
+		if i+1 != numFields || haveObjects {
+			out += "\n"
+		}
+	}
+
+	for i, object := range objects {
+		_, mLength := getDiffString(object.Type)
+		kPrefix := longestMarker - mLength
+		out += formatObjectDiff(object, startPrefix, kPrefix)
+
+		// Avoid a dangling new line
+		if i+1 != numObjects {
+			out += "\n"
+		}
+	}
+
+	return out
+}
+
+// getLongestPrefixes takes a list  of fields and objects and determines the
+// longest field name and the longest marker.
+func getLongestPrefixes(fields []*api.FieldDiff, objects []*api.ObjectDiff) (longestField, longestMarker int) {
+	for _, field := range fields {
+		if l := len(field.Name); l > longestField {
+			longestField = l
+		}
+		if _, l := getDiffString(field.Type); l > longestMarker {
+			longestMarker = l
+		}
+	}
+	for _, obj := range objects {
+		if _, l := getDiffString(obj.Type); l > longestMarker {
+			longestMarker = l
+		}
+	}
+	return longestField, longestMarker
+}
+
+// getDiffString returns a colored diff marker and the length of the string
+// without color annotations.
+func getDiffString(diffType string) (string, int) {
+	switch diffType {
+	case "Added":
+		return "[green]+[reset] ", 2
+	case "Deleted":
+		return "[red]-[reset] ", 2
+	case "Edited":
+		return "[light_yellow]+/-[reset] ", 4
+	default:
+		return "", 0
+	}
+}
+
+// colorAnnotations returns a comma concatonated list of the annotations where
+// the annotations are colored where possible.
+func colorAnnotations(annotations []string) string {
+	l := len(annotations)
+	if l == 0 {
+		return ""
+	}
+
+	colored := make([]string, l)
+	for i, annotation := range annotations {
+		switch annotation {
+		case "forces create":
+			colored[i] = fmt.Sprintf("[green]%s[reset]", annotation)
+		case "forces destroy":
+			colored[i] = fmt.Sprintf("[red]%s[reset]", annotation)
+		case "forces in-place update":
+			colored[i] = fmt.Sprintf("[cyan]%s[reset]", annotation)
+		case "forces create/destroy update":
+			colored[i] = fmt.Sprintf("[yellow]%s[reset]", annotation)
+		default:
+			colored[i] = annotation
+		}
+	}
+
+	return strings.Join(colored, ", ")
+}

command/plan_test.go

@@ -0,0 +1,103 @@
+package command
+
+import (
+	"io/ioutil"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/mitchellh/cli"
+)
+
+func TestPlanCommand_Implements(t *testing.T) {
+	var _ cli.Command = &RunCommand{}
+}
+
+func TestPlanCommand_Fails(t *testing.T) {
+	ui := new(cli.MockUi)
+	cmd := &PlanCommand{Meta: Meta{Ui: ui}}
+
+	// Fails on misuse
+	if code := cmd.Run([]string{"some", "bad", "args"}); code != 1 {
+		t.Fatalf("expected exit code 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, cmd.Help()) {
+		t.Fatalf("expected help output, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fails when specified file does not exist
+	if code := cmd.Run([]string{"/unicorns/leprechauns"}); code != 1 {
+		t.Fatalf("expect exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "Error parsing") {
+		t.Fatalf("expect parsing error, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fails on invalid HCL
+	fh1, err := ioutil.TempFile("", "nomad")
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	defer os.Remove(fh1.Name())
+	if _, err := fh1.WriteString("nope"); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if code := cmd.Run([]string{fh1.Name()}); code != 1 {
+		t.Fatalf("expect exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "Error parsing") {
+		t.Fatalf("expect parsing error, got: %s", err)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fails on invalid job spec
+	fh2, err := ioutil.TempFile("", "nomad")
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	defer os.Remove(fh2.Name())
+	if _, err := fh2.WriteString(`job "job1" {}`); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if code := cmd.Run([]string{fh2.Name()}); code != 1 {
+		t.Fatalf("expect exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "Error validating") {
+		t.Fatalf("expect validation error, got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fails on connection failure (requires a valid job)
+	fh3, err := ioutil.TempFile("", "nomad")
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	defer os.Remove(fh3.Name())
+	_, err = fh3.WriteString(`
+job "job1" {
+	type = "service"
+	datacenters = [ "dc1" ]
+	group "group1" {
+		count = 1
+		task "task1" {
+			driver = "exec"
+			resources = {
+				cpu = 1000
+				disk = 150
+				memory = 512
+			}
+		}
+	}
+}`)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if code := cmd.Run([]string{"-address=nope", fh3.Name()}); code != 1 {
+		t.Fatalf("expected exit code 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "Error during plan") {
+		t.Fatalf("expected failed query error, got: %s", out)
+	}
+}

commands.go

@@ -89,6 +89,13 @@ func Commands(metaPtr *command.Meta) map[string]cli.CommandFactory {
 				Meta: meta,
 			}, nil
 		},
+
+		"plan": func() (cli.Command, error) {
+			return &command.PlanCommand{
+				Meta: meta,
+			}, nil
+		},
+
 		"run": func() (cli.Command, error) {
 			return &command.RunCommand{
 				Meta: meta,

vendor/github.com/mitchellh/colorstring/LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2014 Mitchell Hashimoto
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

vendor/github.com/mitchellh/colorstring/README.md

@@ -0,0 +1,30 @@
+# colorstring [![Build Status](https://travis-ci.org/mitchellh/colorstring.svg)](https://travis-ci.org/mitchellh/colorstring)
+
+colorstring is a [Go](http://www.golang.org) library for outputting colored
+strings to a console using a simple inline syntax in your string to specify
+the color to print as.
+
+For example, the string `[blue]hello [red]world` would output the text
+"hello world" in two colors. The API of colorstring allows for easily disabling
+colors, adding aliases, etc.
+
+## Installation
+
+Standard `go get`:
+
+```
+$ go get github.com/mitchellh/colorstring
+```
+
+## Usage & Example
+
+For usage and examples see the [Godoc](http://godoc.org/github.com/mitchellh/colorstring).
+
+Usage is easy enough:
+
+```go
+colorstring.Println("[blue]Hello [red]World!")
+```
+
+Additionally, the `Colorize` struct can be used to set options such as
+custom colors, color disabling, etc.

vendor/github.com/mitchellh/colorstring/colorstring.go

@@ -0,0 +1,244 @@
+// colorstring provides functions for colorizing strings for terminal
+// output.
+package colorstring
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"regexp"
+	"strings"
+)
+
+// Color colorizes your strings using the default settings.
+//
+// Strings given to Color should use the syntax `[color]` to specify the
+// color for text following. For example: `[blue]Hello` will return "Hello"
+// in blue. See DefaultColors for all the supported colors and attributes.
+//
+// If an unrecognized color is given, it is ignored and assumed to be part
+// of the string. For example: `[hi]world` will result in "[hi]world".
+//
+// A color reset is appended to the end of every string. This will reset
+// the color of following strings when you output this text to the same
+// terminal session.
+//
+// If you want to customize any of this behavior, use the Colorize struct.
+func Color(v string) string {
+	return def.Color(v)
+}
+
+// ColorPrefix returns the color sequence that prefixes the given text.
+//
+// This is useful when wrapping text if you want to inherit the color
+// of the wrapped text. For example, "[green]foo" will return "[green]".
+// If there is no color sequence, then this will return "".
+func ColorPrefix(v string) string {
+	return def.ColorPrefix(v)
+}
+
+// Colorize colorizes your strings, giving you the ability to customize
+// some of the colorization process.
+//
+// The options in Colorize can be set to customize colorization. If you're
+// only interested in the defaults, just use the top Color function directly,
+// which creates a default Colorize.
+type Colorize struct {
+	// Colors maps a color string to the code for that color. The code
+	// is a string so that you can use more complex colors to set foreground,
+	// background, attributes, etc. For example, "boldblue" might be
+	// "1;34"
+	Colors map[string]string
+
+	// If true, color attributes will be ignored. This is useful if you're
+	// outputting to a location that doesn't support colors and you just
+	// want the strings returned.
+	Disable bool
+
+	// Reset, if true, will reset the color after each colorization by
+	// adding a reset code at the end.
+	Reset bool
+}
+
+// Color colorizes a string according to the settings setup in the struct.
+//
+// For more details on the syntax, see the top-level Color function.
+func (c *Colorize) Color(v string) string {
+	matches := parseRe.FindAllStringIndex(v, -1)
+	if len(matches) == 0 {
+		return v
+	}
+
+	result := new(bytes.Buffer)
+	colored := false
+	m := []int{0, 0}
+	for _, nm := range matches {
+		// Write the text in between this match and the last
+		result.WriteString(v[m[1]:nm[0]])
+		m = nm
+
+		var replace string
+		if code, ok := c.Colors[v[m[0]+1:m[1]-1]]; ok {
+			colored = true
+
+			if !c.Disable {
+				replace = fmt.Sprintf("\033[%sm", code)
+			}
+		} else {
+			replace = v[m[0]:m[1]]
+		}
+
+		result.WriteString(replace)
+	}
+	result.WriteString(v[m[1]:])
+
+	if colored && c.Reset && !c.Disable {
+		// Write the clear byte at the end
+		result.WriteString("\033[0m")
+	}
+
+	return result.String()
+}
+
+// ColorPrefix returns the first color sequence that exists in this string.
+//
+// For example: "[green]foo" would return "[green]". If no color sequence
+// exists, then "" is returned. This is especially useful when wrapping
+// colored texts to inherit the color of the wrapped text.
+func (c *Colorize) ColorPrefix(v string) string {
+	return prefixRe.FindString(strings.TrimSpace(v))
+}
+
+// DefaultColors are the default colors used when colorizing.
+//
+// If the color is surrounded in underscores, such as "_blue_", then that
+// color will be used for the background color.
+var DefaultColors map[string]string
+
+func init() {
+	DefaultColors = map[string]string{
+		// Default foreground/background colors
+		"default":   "39",
+		"_default_": "49",
+
+		// Foreground colors
+		"black":         "30",
+		"red":           "31",
+		"green":         "32",
+		"yellow":        "33",
+		"blue":          "34",
+		"magenta":       "35",
+		"cyan":          "36",
+		"light_gray":    "37",
+		"dark_gray":     "90",
+		"light_red":     "91",
+		"light_green":   "92",
+		"light_yellow":  "93",
+		"light_blue":    "94",
+		"light_magenta": "95",
+		"light_cyan":    "96",
+		"white":         "97",
+
+		// Background colors
+		"_black_":         "40",
+		"_red_":           "41",
+		"_green_":         "42",
+		"_yellow_":        "43",
+		"_blue_":          "44",
+		"_magenta_":       "45",
+		"_cyan_":          "46",
+		"_light_gray_":    "47",
+		"_dark_gray_":     "100",
+		"_light_red_":     "101",
+		"_light_green_":   "102",
+		"_light_yellow_":  "103",
+		"_light_blue_":    "104",
+		"_light_magenta_": "105",
+		"_light_cyan_":    "106",
+		"_white_":         "107",
+
+		// Attributes
+		"bold":       "1",
+		"dim":        "2",
+		"underline":  "4",
+		"blink_slow": "5",
+		"blink_fast": "6",
+		"invert":     "7",
+		"hidden":     "8",
+
+		// Reset to reset everything to their defaults
+		"reset":      "0",
+		"reset_bold": "21",
+	}
+
+	def = Colorize{
+		Colors: DefaultColors,
+		Reset:  true,
+	}
+}
+
+var def Colorize
+var parseReRaw = `\[[a-z0-9_-]+\]`
+var parseRe = regexp.MustCompile(`(?i)` + parseReRaw)
+var prefixRe = regexp.MustCompile(`^(?i)(` + parseReRaw + `)+`)
+
+// Print is a convenience wrapper for fmt.Print with support for color codes.
+//
+// Print formats using the default formats for its operands and writes to
+// standard output with support for color codes. Spaces are added between
+// operands when neither is a string. It returns the number of bytes written
+// and any write error encountered.
+func Print(a string) (n int, err error) {
+	return fmt.Print(Color(a))
+}
+
+// Println is a convenience wrapper for fmt.Println with support for color
+// codes.
+//
+// Println formats using the default formats for its operands and writes to
+// standard output with support for color codes. Spaces are always added
+// between operands and a newline is appended. It returns the number of bytes
+// written and any write error encountered.
+func Println(a string) (n int, err error) {
+	return fmt.Println(Color(a))
+}
+
+// Printf is a convenience wrapper for fmt.Printf with support for color codes.
+//
+// Printf formats according to a format specifier and writes to standard output
+// with support for color codes. It returns the number of bytes written and any
+// write error encountered.
+func Printf(format string, a ...interface{}) (n int, err error) {
+	return fmt.Printf(Color(format), a...)
+}
+
+// Fprint is a convenience wrapper for fmt.Fprint with support for color codes.
+//
+// Fprint formats using the default formats for its operands and writes to w
+// with support for color codes. Spaces are added between operands when neither
+// is a string. It returns the number of bytes written and any write error
+// encountered.
+func Fprint(w io.Writer, a string) (n int, err error) {
+	return fmt.Fprint(w, Color(a))
+}
+
+// Fprintln is a convenience wrapper for fmt.Fprintln with support for color
+// codes.
+//
+// Fprintln formats using the default formats for its operands and writes to w
+// with support for color codes. Spaces are always added between operands and a
+// newline is appended. It returns the number of bytes written and any write
+// error encountered.
+func Fprintln(w io.Writer, a string) (n int, err error) {
+	return fmt.Fprintln(w, Color(a))
+}
+
+// Fprintf is a convenience wrapper for fmt.Fprintf with support for color
+// codes.
+//
+// Fprintf formats according to a format specifier and writes to w with support
+// for color codes. It returns the number of bytes written and any write error
+// encountered.
+func Fprintf(w io.Writer, format string, a ...interface{}) (n int, err error) {
+	return fmt.Fprintf(w, Color(format), a...)
+}

vendor/vendor.json

@@ -491,6 +491,12 @@
 			"path": "github.com/mitchellh/cli",
 			"revision": "cb6853d606ea4a12a15ac83cc43503df99fd28fb"
 		},
+		{
+			"checksumSHA1": "ttEN1Aupb7xpPMkQLqb3tzLFdXs=",
+			"path": "github.com/mitchellh/colorstring",
+			"revision": "8631ce90f28644f54aeedcb3e389a85174e067d1",
+			"revisionTime": "2015-09-17T21:48:07Z"
+		},
 		{
 			"path": "github.com/mitchellh/copystructure",
 			"revision": "80adcec1955ee4e97af357c30dee61aadcc02c10"