<commit_before>package colprint

import (
	"io"
	"os"
	"reflect"
	"strconv"
	"fmt"
	"strings"
	"math"
	"sort"
	"github.com/ryanuber/columnize"
)

const TagName = "colprint"

// Config holds configuration used when printing columns
type Config struct {
	// MaxPrintedSliceItems represents the maximum number og slice items to list.
	MaxPrintedSliceItems *int
	// FloatPrecision represents the precision used when printing floats.
	FloatPrecision *int
}

// DefaultPrint prints struct or slice of struct using default config
func DefaultPrint(s interface{}) error {
	return Print(s, nil)
}

// DefaultFprint prints struct or slice to provided Writer using provided config.
func DefaultFprint(w io.Writer, s interface{}) error {
	return Fprint(w, s, nil)
}

// Print prints struct or slice of structs to stdout using provided Config
func Print(s interface{}, c *Config) error {
	return Fprint(os.Stdout, s, c)
}

// Fprint prints struct or slice to provided Writer using provided config.
// If config is nil, default config will be used.
func Fprint(w io.Writer, s interface{}, c *Config) error {
	cp := cPrinter{config: mergeConfig(createDefaultConfig(), c)}
	kind := reflect.TypeOf(s).Kind()
	val := reflect.ValueOf(s)

	// Check if s is a slice/array or not
	if kind == reflect.Slice || kind == reflect.Array {
		// add each item in slice to cPrinter
		for i := 0; i < val.Len(); i ++ {
			if err := cp.add(val.Index(i).Interface()); err != nil {
				return err
			}
		}
	} else {
		// add the item to cPrinter
		if err := cp.add(s); err != nil {
			return err
		}
	}
	// Print to provided Writer
	cp.fprint(w)
	return nil
}

// column represents a column that will be printed by cPrinter
type column struct {
	FieldName string
	Label     string
	Order     int
}

// columns is a sortable list of column structs
type columns []column

func (s columns) Len() int {
	return len(s)
}

func (s columns) Less(i, j int) bool {
	return s[i].Order < s[j].Order
}

func (s columns) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

// cPrinter is the data structure used to print columns
type cPrinter struct {
	// columns (ordered)
	cols columns
	// Map containing values for all columns
	values map[column][]string
	// Keeps track of number of items appended to the ColPrinter
	itemCount int
	// Configuration for the printer
	config *Config
}
// add adds a struct columns and values
func (cp *cPrinter) add(s interface{}) error {
	// Init columns if it's not already done
	if cp.cols == nil {
		cp.init()
		cols, err := cp.findColumns(s)
		if err != nil {
			return err
		}
		cp.cols = cols

		for _, col := range cols {
			cp.initColumn(col)
		}
	}
	// Add values
	for _, col := range cp.cols {
		v := reflect.ValueOf(s)
		val := cp.valueOf(v.FieldByName(col.FieldName).Interface())
		cp.values[col] = append(cp.values[col], val)
	}
	cp.itemCount++
	return nil
}

//fprint prints the columns to the provided io.Writer.
func (cp *cPrinter) fprint(w io.Writer) {
	// Add header line
	str := []string{}
	headers := ""
	for i, col := range cp.cols {
		headers += col.Label
		if i != len(cp.cols)-1 {
			headers += "|"
		}
	}
	str = append(str, headers)

	// Add a line for each item appended
	for i := 0; i < cp.itemCount; i++ {
		vals := ""
		for j, col := range cp.cols {
			vals += cp.values[col][i]
			if j != len(cp.cols)-1 {
				vals += "|"
			}
		}
		str = append(str, vals)
	}
	// Print to given Writer
	fmt.Fprint(w, columnize.SimpleFormat(str)+"\n")
}
// init initializes the array containing columns, and the map containing the values for each column.
func (cp *cPrinter) init() {
	cp.cols = make([]column, 0)
	cp.values = make(map[column][]string)
}
// initColumn initializes the array containing column values.
func (cp *cPrinter) initColumn(col column) {
	cp.values[col] = make([]string, 0)
}

// findColumns extracts which columns should be printed. Returns an error if any field contains a incomplete tag.
func (cp *cPrinter) findColumns(s interface{}) (columns, error) {
	v := reflect.ValueOf(s)
	cols := make(columns, 0)
	for i := 0; i < v.NumField(); i++ {
		field := v.Type().Field(i)
		tag := field.Tag.Get(TagName)

		if tag == "" || tag == "-" {
			continue
		}
		tagVals := strings.Split(tag, ",")

		switch len(tagVals) {
		case 1:
			cols = append(cols, column{field.Name, tagVals[0], math.MaxInt32})
		case 2:
			order, err := strconv.Atoi(tagVals[1])
			if err != nil {
				return nil, fmt.Errorf("Invalid order on field %s", field.Name)
			}
			cols = append(cols, column{field.Name, tagVals[0], order})
		}
	}
	sort.Sort(cols)
	return cols, nil
}

// valueOf returns a string representation of a field.
func (cp *cPrinter) valueOf(i interface{}) string {
	v := reflect.ValueOf(i)
	kind := v.Kind()
	switch kind {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return strconv.FormatInt(v.Int(), 10)
	case reflect.Array, reflect.Slice:
		return cp.valueOfSlice(i)
	case reflect.Bool:
		return strconv.FormatBool(v.Bool())
	case reflect.Float32, reflect.Float64:
		return strconv.FormatFloat(v.Float(), 'f', *cp.config.FloatPrecision, 64)
	case reflect.String:
		return v.String()
	case reflect.Ptr:
		if v.IsNil(){
			return ""
		}
		return cp.valueOf(reflect.Indirect(v).Interface())
	}
	return "<Unsupported kind:"+kind.String()+">"
}
// valueOfSlice returns a string representation of the values in a slice field.
// Returns a maximum of Config.MaxPrintedSliceItems.
func (cp *cPrinter) valueOfSlice(s interface{}) string {
	sliceValue := reflect.ValueOf(s)
	values := ""
	for i := 0; i < sliceValue.Len(); i++ {
		values += cp.valueOf(sliceValue.Index(i).Interface())
		if i == *cp.config.MaxPrintedSliceItems-1 && sliceValue.Len() > *cp.config.MaxPrintedSliceItems {
			values += ",..."
			break
		} else if i < sliceValue.Len()-1 {
			values += ", "
		}
	}
	return values
}
// createDefaultConfig creates a default configuration.
func createDefaultConfig() *Config {
	dMPSI := 3
	dFP := 2
	return &Config{
		MaxPrintedSliceItems: &dMPSI,
		FloatPrecision:       &dFP,
	}
}
// mergeConfig merges the second argument config into the first.
func mergeConfig(a, c *Config) *Config {
	if c != nil {
		if c.MaxPrintedSliceItems != nil {
			*a.MaxPrintedSliceItems = *c.MaxPrintedSliceItems
		}

		if c.FloatPrecision != nil {
			*a.FloatPrecision = *c.FloatPrecision
		}
	}
	return a
}
<commit_msg>Formatted the code<commit_after>package colprint

import (
	"io"
	"os"
	"reflect"
	"strconv"
	"fmt"
	"strings"
	"math"
	"sort"
	"github.com/ryanuber/columnize"
)

const TagName = "colprint"

// Config holds configuration used when printing columns
type Config struct {
	// MaxPrintedSliceItems represents the maximum number og slice items to list.
	MaxPrintedSliceItems *int
	// FloatPrecision represents the precision used when printing floats.
	FloatPrecision *int
}

// DefaultPrint prints struct or slice of struct using default config
func DefaultPrint(s interface{}) error {
	return Print(s, nil)
}

// DefaultFprint prints struct or slice to provided Writer using provided config.
func DefaultFprint(w io.Writer, s interface{}) error {
	return Fprint(w, s, nil)
}

// Print prints struct or slice of structs to stdout using provided Config
func Print(s interface{}, c *Config) error {
	return Fprint(os.Stdout, s, c)
}

// Fprint prints struct or slice to provided Writer using provided config.
// If config is nil, default config will be used.
func Fprint(w io.Writer, s interface{}, c *Config) error {
	cp := cPrinter{config: mergeConfig(createDefaultConfig(), c)}
	kind := reflect.TypeOf(s).Kind()
	val := reflect.ValueOf(s)

	// Check if s is a slice/array or not
	if kind == reflect.Slice || kind == reflect.Array {
		// add each item in slice to cPrinter
		for i := 0; i < val.Len(); i ++ {
			if err := cp.add(val.Index(i).Interface()); err != nil {
				return err
			}
		}
	} else {
		// add the item to cPrinter
		if err := cp.add(s); err != nil {
			return err
		}
	}
	// Print to provided Writer
	cp.fprint(w)
	return nil
}

// column represents a column that will be printed by cPrinter
type column struct {
	FieldName string
	Label     string
	Order     int
}

// columns is a sortable list of column structs
type columns []column

func (s columns) Len() int {
	return len(s)
}

func (s columns) Less(i, j int) bool {
	return s[i].Order < s[j].Order
}

func (s columns) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

// cPrinter is the data structure used to print columns
type cPrinter struct {
	// columns (ordered)
	cols columns
	// Map containing values for all columns
	values map[column][]string
	// Keeps track of number of items appended to the ColPrinter
	itemCount int
	// Configuration for the printer
	config *Config
}

// add adds a struct columns and values
func (cp *cPrinter) add(s interface{}) error {
	// Init columns if it's not already done
	if cp.cols == nil {
		cp.init()
		cols, err := cp.findColumns(s)
		if err != nil {
			return err
		}
		cp.cols = cols

		for _, col := range cols {
			cp.initColumn(col)
		}
	}
	// Add values
	for _, col := range cp.cols {
		v := reflect.ValueOf(s)
		val := cp.valueOf(v.FieldByName(col.FieldName).Interface())
		cp.values[col] = append(cp.values[col], val)
	}
	cp.itemCount++
	return nil
}

//fprint prints the columns to the provided io.Writer.
func (cp *cPrinter) fprint(w io.Writer) {
	// Add header line
	str := []string{}
	headers := ""
	for i, col := range cp.cols {
		headers += col.Label
		if i != len(cp.cols)-1 {
			headers += "|"
		}
	}
	str = append(str, headers)

	// Add a line for each item appended
	for i := 0; i < cp.itemCount; i++ {
		vals := ""
		for j, col := range cp.cols {
			vals += cp.values[col][i]
			if j != len(cp.cols)-1 {
				vals += "|"
			}
		}
		str = append(str, vals)
	}
	// Print to given Writer
	fmt.Fprint(w, columnize.SimpleFormat(str)+"\n")
}

// init initializes the array containing columns, and the map containing the values for each column.
func (cp *cPrinter) init() {
	cp.cols = make([]column, 0)
	cp.values = make(map[column][]string)
}

// initColumn initializes the array containing column values.
func (cp *cPrinter) initColumn(col column) {
	cp.values[col] = make([]string, 0)
}

// findColumns extracts which columns should be printed. Returns an error if any field contains a incomplete tag.
func (cp *cPrinter) findColumns(s interface{}) (columns, error) {
	v := reflect.ValueOf(s)
	cols := make(columns, 0)
	for i := 0; i < v.NumField(); i++ {
		field := v.Type().Field(i)
		tag := field.Tag.Get(TagName)

		if tag == "" || tag == "-" {
			continue
		}
		tagVals := strings.Split(tag, ",")

		switch len(tagVals) {
		case 1:
			cols = append(cols, column{field.Name, tagVals[0], math.MaxInt32})
		case 2:
			order, err := strconv.Atoi(tagVals[1])
			if err != nil {
				return nil, fmt.Errorf("Invalid order on field %s", field.Name)
			}
			cols = append(cols, column{field.Name, tagVals[0], order})
		}
	}
	sort.Sort(cols)
	return cols, nil
}

// valueOf returns a string representation of a field.
func (cp *cPrinter) valueOf(i interface{}) string {
	v := reflect.ValueOf(i)
	kind := v.Kind()
	switch kind {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return strconv.FormatInt(v.Int(), 10)
	case reflect.Array, reflect.Slice:
		return cp.valueOfSlice(i)
	case reflect.Bool:
		return strconv.FormatBool(v.Bool())
	case reflect.Float32, reflect.Float64:
		return strconv.FormatFloat(v.Float(), 'f', *cp.config.FloatPrecision, 64)
	case reflect.String:
		return v.String()
	case reflect.Ptr:
		if v.IsNil() {
			return ""
		}
		return cp.valueOf(reflect.Indirect(v).Interface())
	}
	return "<Unsupported kind:" + kind.String() + ">"
}

// valueOfSlice returns a string representation of the values in a slice field.
// Returns a maximum of Config.MaxPrintedSliceItems.
func (cp *cPrinter) valueOfSlice(s interface{}) string {
	sliceValue := reflect.ValueOf(s)
	values := ""
	for i := 0; i < sliceValue.Len(); i++ {
		values += cp.valueOf(sliceValue.Index(i).Interface())
		if i == *cp.config.MaxPrintedSliceItems-1 && sliceValue.Len() > *cp.config.MaxPrintedSliceItems {
			values += ",..."
			break
		} else if i < sliceValue.Len()-1 {
			values += ", "
		}
	}
	return values
}

// createDefaultConfig creates a default configuration.
func createDefaultConfig() *Config {
	dMPSI := 3
	dFP := 2
	return &Config{
		MaxPrintedSliceItems: &dMPSI,
		FloatPrecision:       &dFP,
	}
}

// mergeConfig merges the second argument config into the first.
func mergeConfig(a, c *Config) *Config {
	if c != nil {
		if c.MaxPrintedSliceItems != nil {
			*a.MaxPrintedSliceItems = *c.MaxPrintedSliceItems
		}

		if c.FloatPrecision != nil {
			*a.FloatPrecision = *c.FloatPrecision
		}
	}
	return a
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go-Commander Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Based on the original work by The Go Authors:
// Copyright 2011 The Go Authors.  All rights reserved.

// commander helps creating command line programs whose arguments are flags,
// commands and subcommands.
package commander

import (
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"

	"github.com/sbinet/go-flag"
)

// A Commander holds the configuration for the command line tool.
type Commander struct {
	// Name is the command name, usually the executable's name.
	Name string
	// Commands is the list of commands supported by this commander program.
	Commands []*Command
	// Flag is a set of flags for the whole commander. It should not be
	// changed after Run() is called.
	Flag *flag.FlagSet
}

// Run executes the commander using the provided arguments. The command
// matching the first argument is executed and it receives the remaining
// arguments.
func (c *Commander) Run(args []string) error {
	if c == nil {
		return fmt.Errorf("Called Run() on a nil Commander")
	}
	if c.Flag == nil {
		c.Flag = flag.NewFlagSet(c.Name, flag.ExitOnError)
	}
	if c.Flag.Usage == nil {
		c.Flag.Usage = func() {
			if err := c.usage(); err != nil {
				fmt.Fprintln(os.Stderr, err)
			}
		}
	}
	if !c.Flag.Parsed() {
		if err := c.Flag.Parse(args); err != nil {
			return fmt.Errorf("Commander.Main flag parsing failure: %v", err)
		}
	}
	if len(args) < 1 {
		if err := c.usage(); err != nil {
			return err
		}
		return fmt.Errorf("Not enough arguments provided")
	}

	if args[0] == "help" {
		return c.help(args[1:])
	}

	for _, cmd := range c.Commands {
		if cmd.Name() == args[0] && cmd.Runnable() {
			cmd.Flag.Usage = func() { cmd.Usage() }
			if cmd.CustomFlags {
				args = args[1:]
			} else {
				cmd.Flag.Parse(args[1:])
				args = cmd.Flag.Args()
			}
			cmd.Run(cmd, args)
			return nil
		}
	}

	return fmt.Errorf("unknown subcommand %q\nRun 'help' for usage.\n", args[0])
}

func (c *Commander) usage() error {
	err := tmpl(os.Stderr, usageTemplate, c)
	if err != nil {
		fmt.Println(err)
	}
	return err
}

// help implements the 'help' command.
func (c *Commander) help(args []string) error {
	if len(args) == 0 {
		return c.usage()
	}
	if len(args) != 1 {
		return fmt.Errorf("usage: %v help command\n\nToo many arguments given.\n", c.Name)
	}

	arg := args[0]

	for _, cmd := range c.Commands {
		if cmd.Name() == arg {
			c := struct {
				*Command
				ProgramName string
			}{cmd, c.Name}
			return tmpl(os.Stdout, helpTemplate, c)
		}
	}

	return fmt.Errorf("Unknown help topic %#q.  Run '%v help'.\n", arg, c.Name)
}

// A Command is an implementation of a subcommand.
type Command struct {
	// Run runs the command.
	// The args are the arguments after the command name.
	Run func(cmd *Command, args []string)

	// UsageLine is the one-line usage message.
	// The first word in the line is taken to be the command name.
	UsageLine string

	// Short is the short description shown in the 'help' output.
	Short string

	// Long is the long message shown in the 'help <this-command>' output.
	Long string

	// Flag is a set of flags specific to this command.
	Flag flag.FlagSet

	// CustomFlags indicates that the command will do its own
	// flag parsing.
	CustomFlags bool
}

// Name returns the command's name: the first word in the usage line.
func (c *Command) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Usage prints the usage details to the standard error output.
func (c *Command) Usage() {
	fmt.Fprintf(os.Stderr, "usage: %s\n\n", c.UsageLine)
	fmt.Fprintf(os.Stderr, "%s\n", strings.TrimSpace(c.Long))
}

// Runnable reports whether the command can be run; otherwise
// it is a documentation pseudo-command such as importpath.
func (c *Command) Runnable() bool {
	return c.Run != nil
}

var usageTemplate = `Usage:

	{{.Name}} command [arguments]

The commands are:
{{range .Commands}}{{if .Runnable}}
    {{.Name | printf "%-11s"}} {{.Short}}{{end}}{{end}}

Use "{{$.Name}} help [command]" for more information about a command.

Additional help topics:
{{range .Commands}}{{if not .Runnable}}
    {{.Name | printf "%-11s"}} {{.Short}}{{end}}{{end}}

Use "{{.Name}} help [topic]" for more information about that topic.

`

var helpTemplate = `{{if .Runnable}}Usage: {{.ProgramName}} {{.UsageLine}}

{{end}}{{.Long | trim}}
`

// tmpl executes the given template text on data, writing the result to w.
func tmpl(w io.Writer, text string, data interface{}) error {
	t := template.New("top")
	t.Funcs(template.FuncMap{"trim": strings.TrimSpace})
	template.Must(t.Parse(text))
	return t.Execute(w, data)
}
<commit_msg>automatically display option-flags for subcommands with 'help'<commit_after>// Copyright 2012 The Go-Commander Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Based on the original work by The Go Authors:
// Copyright 2011 The Go Authors.  All rights reserved.

// commander helps creating command line programs whose arguments are flags,
// commands and subcommands.
package commander

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"
	"text/template"

	"github.com/sbinet/go-flag"
)

// A Commander holds the configuration for the command line tool.
type Commander struct {
	// Name is the command name, usually the executable's name.
	Name string
	// Commands is the list of commands supported by this commander program.
	Commands []*Command
	// Flag is a set of flags for the whole commander. It should not be
	// changed after Run() is called.
	Flag *flag.FlagSet
}

// Run executes the commander using the provided arguments. The command
// matching the first argument is executed and it receives the remaining
// arguments.
func (c *Commander) Run(args []string) error {
	if c == nil {
		return fmt.Errorf("Called Run() on a nil Commander")
	}
	if c.Flag == nil {
		c.Flag = flag.NewFlagSet(c.Name, flag.ExitOnError)
	}
	if c.Flag.Usage == nil {
		c.Flag.Usage = func() {
			if err := c.usage(); err != nil {
				fmt.Fprintln(os.Stderr, err)
			}
		}
	}
	if !c.Flag.Parsed() {
		if err := c.Flag.Parse(args); err != nil {
			return fmt.Errorf("Commander.Main flag parsing failure: %v", err)
		}
	}
	if len(args) < 1 {
		if err := c.usage(); err != nil {
			return err
		}
		return fmt.Errorf("Not enough arguments provided")
	}

	if args[0] == "help" {
		return c.help(args[1:])
	}

	for _, cmd := range c.Commands {
		if cmd.Name() == args[0] && cmd.Runnable() {
			cmd.Flag.Usage = func() { cmd.Usage() }
			if cmd.CustomFlags {
				args = args[1:]
			} else {
				cmd.Flag.Parse(args[1:])
				args = cmd.Flag.Args()
			}
			cmd.Run(cmd, args)
			return nil
		}
	}

	return fmt.Errorf("unknown subcommand %q\nRun 'help' for usage.\n", args[0])
}

func (c *Commander) usage() error {
	err := tmpl(os.Stderr, usageTemplate, c)
	if err != nil {
		fmt.Println(err)
	}
	return err
}

// help implements the 'help' command.
func (c *Commander) help(args []string) error {
	if len(args) == 0 {
		return c.usage()
	}
	if len(args) != 1 {
		return fmt.Errorf("usage: %v help command\n\nToo many arguments given.\n", c.Name)
	}

	arg := args[0]

	for _, cmd := range c.Commands {
		if cmd.Name() == arg {
			c := struct {
				*Command
				ProgramName string
			}{cmd, c.Name}
			return tmpl(os.Stdout, helpTemplate, c)
		}
	}

	return fmt.Errorf("Unknown help topic %#q.  Run '%v help'.\n", arg, c.Name)
}

// A Command is an implementation of a subcommand.
type Command struct {
	// Run runs the command.
	// The args are the arguments after the command name.
	Run func(cmd *Command, args []string)

	// UsageLine is the one-line usage message.
	// The first word in the line is taken to be the command name.
	UsageLine string

	// Short is the short description shown in the 'help' output.
	Short string

	// Long is the long message shown in the 'help <this-command>' output.
	Long string

	// Flag is a set of flags specific to this command.
	Flag flag.FlagSet

	// CustomFlags indicates that the command will do its own
	// flag parsing.
	CustomFlags bool
}

// Name returns the command's name: the first word in the usage line.
func (c *Command) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Usage prints the usage details to the standard error output.
func (c *Command) Usage() {
	fmt.Fprintf(os.Stderr, "usage: %s\n\n", c.UsageLine)
	fmt.Fprintf(os.Stderr, "%s\n", strings.TrimSpace(c.Long))
}

// flagUsage returns the usage details as a string
func flagUsage(flagset *flag.FlagSet) string {
	if flagset == nil {
		return ""
	}
	var buf bytes.Buffer
	flagset.SetOutput(&buf)
	fmt.Fprintf(&buf, "\noptions:\n")
	if flagset.Usage != nil {
		flagset.Usage()
	} else {
		flagset.PrintDefaults()
	}
	return string(buf.Bytes())
}

// Runnable reports whether the command can be run; otherwise
// it is a documentation pseudo-command such as importpath.
func (c *Command) Runnable() bool {
	return c.Run != nil
}

var usageTemplate = `Usage:

	{{.Name}} command [arguments]

The commands are:
{{range .Commands}}{{if .Runnable}}
    {{.Name | printf "%-11s"}} {{.Short}}{{end}}{{end}}

Use "{{$.Name}} help [command]" for more information about a command.

Additional help topics:
{{range .Commands}}{{if not .Runnable}}
    {{.Name | printf "%-11s"}} {{.Short}}{{end}}{{end}}

Use "{{.Name}} help [topic]" for more information about that topic.

`

var helpTemplate = `{{if .Runnable}}Usage: {{.ProgramName}} {{.UsageLine}}

{{end}}{{.Long | trim}}
{{.Flag | flagUsage}}
`

// tmpl executes the given template text on data, writing the result to w.
func tmpl(w io.Writer, text string, data interface{}) error {
	t := template.New("top")
	t.Funcs(template.FuncMap{
		"trim": strings.TrimSpace,
		"flagUsage": flagUsage,
	})
	template.Must(t.Parse(text))
	return t.Execute(w, data)
}
<|endoftext|>


<commit_before>// Copyright 2010 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// MockGen generates mock implementations of Go interfaces.
package main

// TODO: This does not support recursive embedded interfaces.
// TODO: This does not support embedding package-local interfaces in a separate file.

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"go/token"
	"io"
	"log"
	"os"
	"path"
	"strconv"
	"strings"
	"unicode"

	"github.com/golang/mock/mockgen/model"
)

const (
	gomockImportPath = "github.com/golang/mock/gomock"
)

var (
	source      = flag.String("source", "", "(source mode) Input Go source file; enables source mode.")
	destination = flag.String("destination", "", "Output file; defaults to stdout.")
	packageOut  = flag.String("package", "", "Package of the generated code; defaults to the package of the input with a 'mock_' prefix.")
	selfPackage = flag.String("self_package", "", "If set, the package this mock will be part of.")

	debugParser = flag.Bool("debug_parser", false, "Print out parser results only.")
)

func main() {
	flag.Usage = usage
	flag.Parse()

	var pkg *model.Package
	var err error
	if *source != "" {
		pkg, err = ParseFile(*source)
	} else {
		if flag.NArg() != 2 {
			log.Fatal("Expected exactly two arguments")
		}
		pkg, err = Reflect(flag.Arg(0), strings.Split(flag.Arg(1), ","))
	}
	if err != nil {
		log.Fatalf("Loading input failed: %v", err)
	}

	if *debugParser {
		pkg.Print(os.Stdout)
		return
	}

	dst := os.Stdout
	if len(*destination) > 0 {
		f, err := os.Create(*destination)
		if err != nil {
			log.Fatalf("Failed opening destination file: %v", err)
		}
		defer f.Close()
		dst = f
	}

	packageName := *packageOut
	if packageName == "" {
		// pkg.Name in reflect mode is the base name of the import path,
		// which might have characters that are illegal to have in package names.
		packageName = "mock_" + sanitize(pkg.Name)
	}

	g := new(generator)
	if *source != "" {
		g.filename = *source
	} else {
		g.srcPackage = flag.Arg(0)
		g.srcInterfaces = flag.Arg(1)
	}
	if err := g.Generate(pkg, packageName); err != nil {
		log.Fatalf("Failed generating mock: %v", err)
	}
	if _, err := dst.Write(g.Output()); err != nil {
		log.Fatalf("Failed writing to destination: %v", err)
	}
}

func usage() {
	io.WriteString(os.Stderr, usageText)
	flag.PrintDefaults()
}

const usageText = `mockgen has two modes of operation: source and reflect.

Source mode generates mock interfaces from a source file.
It is enabled by using the -source flag. Other flags that
may be useful in this mode are -imports and -aux_files.
Example:
	mockgen -source=foo.go [other options]

Reflect mode generates mock interfaces by building a program
that uses reflection to understand interfaces. It is enabled
by passing two non-flag arguments: an import path, and a
comma-separated list of symbols.
Example:
	mockgen database/sql/driver Conn,Driver

`

type generator struct {
	buf    bytes.Buffer
	indent string

	filename                  string // may be empty
	srcPackage, srcInterfaces string // may be empty

	packageMap map[string]string // map from import path to package name
}

func (g *generator) p(format string, args ...interface{}) {
	fmt.Fprintf(&g.buf, g.indent+format+"\n", args...)
}

func (g *generator) in() {
	g.indent += "\t"
}

func (g *generator) out() {
	if len(g.indent) > 0 {
		g.indent = g.indent[0 : len(g.indent)-1]
	}
}

func removeDot(s string) string {
	if len(s) > 0 && s[len(s)-1] == '.' {
		return s[0 : len(s)-1]
	}
	return s
}

// sanitize cleans up a string to make a suitable package name.
func sanitize(s string) string {
	t := ""
	for _, r := range s {
		if t == "" {
			if unicode.IsLetter(r) || r == '_' {
				t += string(r)
				continue
			}
		} else {
			if unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' {
				t += string(r)
				continue
			}
		}
		t += "_"
	}
	if t == "_" {
		t = "x"
	}
	return t
}

func (g *generator) Generate(pkg *model.Package, pkgName string) error {
	g.p("// Automatically generated by MockGen. DO NOT EDIT!")
	if g.filename != "" {
		g.p("// Source: %v", g.filename)
	} else {
		g.p("// Source: %v (interfaces: %v)", g.srcPackage, g.srcInterfaces)
	}
	g.p("")

	// Get all required imports, and generate unique names for them all.
	im := pkg.Imports()
	im[gomockImportPath] = true
	g.packageMap = make(map[string]string, len(im))
	localNames := make(map[string]bool, len(im))
	for pth := range im {
		base := sanitize(path.Base(pth))

		// Local names for an imported package can usually be the basename of the import path.
		// A couple of situations don't permit that, such as duplicate local names
		// (e.g. importing "html/template" and "text/template"), or where the basename is
		// a keyword (e.g. "foo/case").
		// try base0, base1, ...
		pkgName := base
		i := 0
		for localNames[pkgName] || token.Lookup(pkgName).IsKeyword() {
			pkgName = base + strconv.Itoa(i)
			i++
		}

		g.packageMap[pth] = pkgName
		localNames[pkgName] = true
	}

	g.p("package %v", pkgName)
	g.p("")
	g.p("import (")
	g.in()
	for path, pkg := range g.packageMap {
		if path == *selfPackage {
			continue
		}
		g.p("%v %q", pkg, path)
	}
	for _, path := range pkg.DotImports {
		g.p(". %q", path)
	}
	g.out()
	g.p(")")

	for _, intf := range pkg.Interfaces {
		if err := g.GenerateMockInterface(intf); err != nil {
			return err
		}
	}

	return nil
}

// The name of the mock type to use for the given interface identifier.
func mockName(typeName string) string {
	return "Mock" + typeName
}

func (g *generator) GenerateMockInterface(intf *model.Interface) error {
	mockType := mockName(intf.Name)

	g.p("")
	g.p("// Mock of %v interface", intf.Name)
	g.p("type %v struct {", mockType)
	g.in()
	g.p("ctrl     *gomock.Controller")
	g.p("recorder *_%vRecorder", mockType)
	g.out()
	g.p("}")
	g.p("")

	g.p("// Recorder for %v (not exported)", mockType)
	g.p("type _%vRecorder struct {", mockType)
	g.in()
	g.p("mock *%v", mockType)
	g.out()
	g.p("}")
	g.p("")

	// TODO: Re-enable this if we can import the interface reliably.
	//g.p("// Verify that the mock satisfies the interface at compile time.")
	//g.p("var _ %v = (*%v)(nil)", typeName, mockType)
	//g.p("")

	g.p("func New%v(ctrl *gomock.Controller) *%v {", mockType, mockType)
	g.in()
	g.p("mock := &%v{ctrl: ctrl}", mockType)
	g.p("mock.recorder = &_%vRecorder{mock}", mockType)
	g.p("return mock")
	g.out()
	g.p("}")
	g.p("")

	// XXX: possible name collision here if someone has EXPECT in their interface.
	g.p("func (_m *%v) EXPECT() *_%vRecorder {", mockType, mockType)
	g.in()
	g.p("return _m.recorder")
	g.out()
	g.p("}")

	g.GenerateMockMethods(mockType, intf, *selfPackage)

	return nil
}

func (g *generator) GenerateMockMethods(mockType string, intf *model.Interface, pkgOverride string) {
	for _, m := range intf.Methods {
		g.p("")
		g.GenerateMockMethod(mockType, m, pkgOverride)
		g.p("")
		g.GenerateMockRecorderMethod(mockType, m)
	}
}

// GenerateMockMethod generates a mock method implementation.
// If non-empty, pkgOverride is the package in which unqualified types reside.
func (g *generator) GenerateMockMethod(mockType string, m *model.Method, pkgOverride string) error {
	args := make([]string, len(m.In))
	argNames := make([]string, len(m.In))
	for i, p := range m.In {
		name := p.Name
		if name == "" {
			name = fmt.Sprintf("_param%d", i)
		}
		ts := p.Type.String(g.packageMap, pkgOverride)
		args[i] = name + " " + ts
		argNames[i] = name
	}
	if m.Variadic != nil {
		name := m.Variadic.Name
		if name == "" {
			name = fmt.Sprintf("_param%d", len(m.In))
		}
		ts := m.Variadic.Type.String(g.packageMap, pkgOverride)
		args = append(args, name+" ..."+ts)
		argNames = append(argNames, name)
	}
	argString := strings.Join(args, ", ")

	rets := make([]string, len(m.Out))
	for i, p := range m.Out {
		rets[i] = p.Type.String(g.packageMap, pkgOverride)
	}
	retString := strings.Join(rets, ", ")
	if len(rets) > 1 {
		retString = "(" + retString + ")"
	}
	if retString != "" {
		retString = " " + retString
	}

	g.p("func (_m *%v) %v(%v)%v {", mockType, m.Name, argString, retString)
	g.in()

	callArgs := strings.Join(argNames, ", ")
	if callArgs != "" {
		callArgs = ", " + callArgs
	}
	if m.Variadic != nil {
		// Non-trivial. The generated code must build a []interface{},
		// but the variadic argument may be any type.
		g.p("_s := []interface{}{%s}", strings.Join(argNames[:len(argNames)-1], ", "))
		g.p("for _, _x := range %s {", argNames[len(argNames)-1])
		g.in()
		g.p("_s = append(_s, _x)")
		g.out()
		g.p("}")
		callArgs = ", _s..."
	}
	if len(m.Out) == 0 {
		g.p(`_m.ctrl.Call(_m, "%v"%v)`, m.Name, callArgs)
	} else {
		g.p(`ret := _m.ctrl.Call(_m, "%v"%v)`, m.Name, callArgs)

		// Go does not allow "naked" type assertions on nil values, so we use the two-value form here.
		// The value of that is either (x.(T), true) or (Z, false), where Z is the zero value for T.
		// Happily, this coincides with the semantics we want here.
		retNames := make([]string, len(rets))
		for i, t := range rets {
			retNames[i] = fmt.Sprintf("ret%d", i)
			g.p("%s, _ := ret[%d].(%s)", retNames[i], i, t)
		}
		g.p("return " + strings.Join(retNames, ", "))
	}

	g.out()
	g.p("}")
	return nil
}

func (g *generator) GenerateMockRecorderMethod(mockType string, m *model.Method) error {
	nargs := len(m.In)
	args := make([]string, nargs)
	for i := 0; i < nargs; i++ {
		args[i] = "arg" + strconv.Itoa(i)
	}
	argString := strings.Join(args, ", ")
	if nargs > 0 {
		argString += " interface{}"
	}
	if m.Variadic != nil {
		if nargs > 0 {
			argString += ", "
		}
		argString += fmt.Sprintf("arg%d ...interface{}", nargs)
	}

	g.p("func (_mr *_%vRecorder) %v(%v) *gomock.Call {", mockType, m.Name, argString)
	g.in()

	callArgs := strings.Join(args, ", ")
	if nargs > 0 {
		callArgs = ", " + callArgs
	}
	if m.Variadic != nil {
		if nargs == 0 {
			// Easy: just use ... to push the arguments through.
			callArgs = ", arg0..."
		} else {
			// Hard: create a temporary slice.
			g.p("_s := append([]interface{}{%s}, arg%d...)", strings.Join(args, ", "), nargs)
			callArgs = ", _s..."
		}
	}
	g.p(`return _mr.mock.ctrl.RecordCall(_mr.mock, "%v"%v)`, m.Name, callArgs)

	g.out()
	g.p("}")
	return nil
}

// Output returns the generator's output, formatted in the standard Go style.
func (g *generator) Output() []byte {
	src, err := format.Source(g.buf.Bytes())
	if err != nil {
		log.Fatalf("Failed to format generated source code: %s\n%s", err, g.buf.String())
	}
	return src
}
<commit_msg>add import_prefix to specify a vendored gomock pkg<commit_after>// Copyright 2010 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// MockGen generates mock implementations of Go interfaces.
package main

// TODO: This does not support recursive embedded interfaces.
// TODO: This does not support embedding package-local interfaces in a separate file.

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"go/token"
	"io"
	"log"
	"os"
	"path"
	"strconv"
	"strings"
	"unicode"

	"github.com/golang/mock/mockgen/model"
)

const (
	gomockImportPath = "github.com/golang/mock/gomock"
)

var (
	source       = flag.String("source", "", "(source mode) Input Go source file; enables source mode.")
	destination  = flag.String("destination", "", "Output file; defaults to stdout.")
	packageOut   = flag.String("package", "", "Package of the generated code; defaults to the package of the input with a 'mock_' prefix.")
	selfPackage  = flag.String("self_package", "", "If set, the package this mock will be part of.")
	importPrefix = flag.String("import_prefix", "", "If set, the prefix for the import of the gomock package")

	debugParser = flag.Bool("debug_parser", false, "Print out parser results only.")
)

func main() {
	flag.Usage = usage
	flag.Parse()

	var pkg *model.Package
	var err error
	if *source != "" {
		pkg, err = ParseFile(*source)
	} else {
		if flag.NArg() != 2 {
			log.Fatal("Expected exactly two arguments")
		}
		pkg, err = Reflect(flag.Arg(0), strings.Split(flag.Arg(1), ","))
	}
	if err != nil {
		log.Fatalf("Loading input failed: %v", err)
	}

	if *debugParser {
		pkg.Print(os.Stdout)
		return
	}

	dst := os.Stdout
	if len(*destination) > 0 {
		f, err := os.Create(*destination)
		if err != nil {
			log.Fatalf("Failed opening destination file: %v", err)
		}
		defer f.Close()
		dst = f
	}

	packageName := *packageOut
	if packageName == "" {
		// pkg.Name in reflect mode is the base name of the import path,
		// which might have characters that are illegal to have in package names.
		packageName = "mock_" + sanitize(pkg.Name)
	}

	g := new(generator)
	if *source != "" {
		g.filename = *source
	} else {
		g.srcPackage = flag.Arg(0)
		g.srcInterfaces = flag.Arg(1)
	}
	if err := g.Generate(pkg, packageName); err != nil {
		log.Fatalf("Failed generating mock: %v", err)
	}
	if _, err := dst.Write(g.Output()); err != nil {
		log.Fatalf("Failed writing to destination: %v", err)
	}
}

func usage() {
	io.WriteString(os.Stderr, usageText)
	flag.PrintDefaults()
}

const usageText = `mockgen has two modes of operation: source and reflect.

Source mode generates mock interfaces from a source file.
It is enabled by using the -source flag. Other flags that
may be useful in this mode are -imports and -aux_files.
Example:
	mockgen -source=foo.go [other options]

Reflect mode generates mock interfaces by building a program
that uses reflection to understand interfaces. It is enabled
by passing two non-flag arguments: an import path, and a
comma-separated list of symbols.
Example:
	mockgen database/sql/driver Conn,Driver

`

type generator struct {
	buf    bytes.Buffer
	indent string

	filename                  string // may be empty
	srcPackage, srcInterfaces string // may be empty

	packageMap map[string]string // map from import path to package name
}

func (g *generator) p(format string, args ...interface{}) {
	fmt.Fprintf(&g.buf, g.indent+format+"\n", args...)
}

func (g *generator) in() {
	g.indent += "\t"
}

func (g *generator) out() {
	if len(g.indent) > 0 {
		g.indent = g.indent[0 : len(g.indent)-1]
	}
}

func removeDot(s string) string {
	if len(s) > 0 && s[len(s)-1] == '.' {
		return s[0 : len(s)-1]
	}
	return s
}

// sanitize cleans up a string to make a suitable package name.
func sanitize(s string) string {
	t := ""
	for _, r := range s {
		if t == "" {
			if unicode.IsLetter(r) || r == '_' {
				t += string(r)
				continue
			}
		} else {
			if unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' {
				t += string(r)
				continue
			}
		}
		t += "_"
	}
	if t == "_" {
		t = "x"
	}
	return t
}

func (g *generator) Generate(pkg *model.Package, pkgName string) error {
	g.p("// Automatically generated by MockGen. DO NOT EDIT!")
	if g.filename != "" {
		g.p("// Source: %v", g.filename)
	} else {
		g.p("// Source: %v (interfaces: %v)", g.srcPackage, g.srcInterfaces)
	}
	g.p("")

	// Get all required imports, and generate unique names for them all.
	im := pkg.Imports()
	gomockPath := path.Join(*importPrefix, gomockImportPath)
	im[gomockPath] = true
	g.packageMap = make(map[string]string, len(im))
	localNames := make(map[string]bool, len(im))
	for pth := range im {
		base := sanitize(path.Base(pth))

		// Local names for an imported package can usually be the basename of the import path.
		// A couple of situations don't permit that, such as duplicate local names
		// (e.g. importing "html/template" and "text/template"), or where the basename is
		// a keyword (e.g. "foo/case").
		// try base0, base1, ...
		pkgName := base
		i := 0
		for localNames[pkgName] || token.Lookup(pkgName).IsKeyword() {
			pkgName = base + strconv.Itoa(i)
			i++
		}

		g.packageMap[pth] = pkgName
		localNames[pkgName] = true
	}

	g.p("package %v", pkgName)
	g.p("")
	g.p("import (")
	g.in()
	for path, pkg := range g.packageMap {
		if path == *selfPackage {
			continue
		}
		g.p("%v %q", pkg, path)
	}
	for _, path := range pkg.DotImports {
		g.p(". %q", path)
	}
	g.out()
	g.p(")")

	for _, intf := range pkg.Interfaces {
		if err := g.GenerateMockInterface(intf); err != nil {
			return err
		}
	}

	return nil
}

// The name of the mock type to use for the given interface identifier.
func mockName(typeName string) string {
	return "Mock" + typeName
}

func (g *generator) GenerateMockInterface(intf *model.Interface) error {
	mockType := mockName(intf.Name)

	g.p("")
	g.p("// Mock of %v interface", intf.Name)
	g.p("type %v struct {", mockType)
	g.in()
	g.p("ctrl     *gomock.Controller")
	g.p("recorder *_%vRecorder", mockType)
	g.out()
	g.p("}")
	g.p("")

	g.p("// Recorder for %v (not exported)", mockType)
	g.p("type _%vRecorder struct {", mockType)
	g.in()
	g.p("mock *%v", mockType)
	g.out()
	g.p("}")
	g.p("")

	// TODO: Re-enable this if we can import the interface reliably.
	//g.p("// Verify that the mock satisfies the interface at compile time.")
	//g.p("var _ %v = (*%v)(nil)", typeName, mockType)
	//g.p("")

	g.p("func New%v(ctrl *gomock.Controller) *%v {", mockType, mockType)
	g.in()
	g.p("mock := &%v{ctrl: ctrl}", mockType)
	g.p("mock.recorder = &_%vRecorder{mock}", mockType)
	g.p("return mock")
	g.out()
	g.p("}")
	g.p("")

	// XXX: possible name collision here if someone has EXPECT in their interface.
	g.p("func (_m *%v) EXPECT() *_%vRecorder {", mockType, mockType)
	g.in()
	g.p("return _m.recorder")
	g.out()
	g.p("}")

	g.GenerateMockMethods(mockType, intf, *selfPackage)

	return nil
}

func (g *generator) GenerateMockMethods(mockType string, intf *model.Interface, pkgOverride string) {
	for _, m := range intf.Methods {
		g.p("")
		g.GenerateMockMethod(mockType, m, pkgOverride)
		g.p("")
		g.GenerateMockRecorderMethod(mockType, m)
	}
}

// GenerateMockMethod generates a mock method implementation.
// If non-empty, pkgOverride is the package in which unqualified types reside.
func (g *generator) GenerateMockMethod(mockType string, m *model.Method, pkgOverride string) error {
	args := make([]string, len(m.In))
	argNames := make([]string, len(m.In))
	for i, p := range m.In {
		name := p.Name
		if name == "" {
			name = fmt.Sprintf("_param%d", i)
		}
		ts := p.Type.String(g.packageMap, pkgOverride)
		args[i] = name + " " + ts
		argNames[i] = name
	}
	if m.Variadic != nil {
		name := m.Variadic.Name
		if name == "" {
			name = fmt.Sprintf("_param%d", len(m.In))
		}
		ts := m.Variadic.Type.String(g.packageMap, pkgOverride)
		args = append(args, name+" ..."+ts)
		argNames = append(argNames, name)
	}
	argString := strings.Join(args, ", ")

	rets := make([]string, len(m.Out))
	for i, p := range m.Out {
		rets[i] = p.Type.String(g.packageMap, pkgOverride)
	}
	retString := strings.Join(rets, ", ")
	if len(rets) > 1 {
		retString = "(" + retString + ")"
	}
	if retString != "" {
		retString = " " + retString
	}

	g.p("func (_m *%v) %v(%v)%v {", mockType, m.Name, argString, retString)
	g.in()

	callArgs := strings.Join(argNames, ", ")
	if callArgs != "" {
		callArgs = ", " + callArgs
	}
	if m.Variadic != nil {
		// Non-trivial. The generated code must build a []interface{},
		// but the variadic argument may be any type.
		g.p("_s := []interface{}{%s}", strings.Join(argNames[:len(argNames)-1], ", "))
		g.p("for _, _x := range %s {", argNames[len(argNames)-1])
		g.in()
		g.p("_s = append(_s, _x)")
		g.out()
		g.p("}")
		callArgs = ", _s..."
	}
	if len(m.Out) == 0 {
		g.p(`_m.ctrl.Call(_m, "%v"%v)`, m.Name, callArgs)
	} else {
		g.p(`ret := _m.ctrl.Call(_m, "%v"%v)`, m.Name, callArgs)

		// Go does not allow "naked" type assertions on nil values, so we use the two-value form here.
		// The value of that is either (x.(T), true) or (Z, false), where Z is the zero value for T.
		// Happily, this coincides with the semantics we want here.
		retNames := make([]string, len(rets))
		for i, t := range rets {
			retNames[i] = fmt.Sprintf("ret%d", i)
			g.p("%s, _ := ret[%d].(%s)", retNames[i], i, t)
		}
		g.p("return " + strings.Join(retNames, ", "))
	}

	g.out()
	g.p("}")
	return nil
}

func (g *generator) GenerateMockRecorderMethod(mockType string, m *model.Method) error {
	nargs := len(m.In)
	args := make([]string, nargs)
	for i := 0; i < nargs; i++ {
		args[i] = "arg" + strconv.Itoa(i)
	}
	argString := strings.Join(args, ", ")
	if nargs > 0 {
		argString += " interface{}"
	}
	if m.Variadic != nil {
		if nargs > 0 {
			argString += ", "
		}
		argString += fmt.Sprintf("arg%d ...interface{}", nargs)
	}

	g.p("func (_mr *_%vRecorder) %v(%v) *gomock.Call {", mockType, m.Name, argString)
	g.in()

	callArgs := strings.Join(args, ", ")
	if nargs > 0 {
		callArgs = ", " + callArgs
	}
	if m.Variadic != nil {
		if nargs == 0 {
			// Easy: just use ... to push the arguments through.
			callArgs = ", arg0..."
		} else {
			// Hard: create a temporary slice.
			g.p("_s := append([]interface{}{%s}, arg%d...)", strings.Join(args, ", "), nargs)
			callArgs = ", _s..."
		}
	}
	g.p(`return _mr.mock.ctrl.RecordCall(_mr.mock, "%v"%v)`, m.Name, callArgs)

	g.out()
	g.p("}")
	return nil
}

// Output returns the generator's output, formatted in the standard Go style.
func (g *generator) Output() []byte {
	src, err := format.Source(g.buf.Bytes())
	if err != nil {
		log.Fatalf("Failed to format generated source code: %s\n%s", err, g.buf.String())
	}
	return src
}
<|endoftext|>


<commit_before>// Copyright 2012 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

// This file contains the model construction by reflection.

import (
	"bytes"
	"encoding/gob"
	"flag"
	"go/build"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"text/template"

	"github.com/golang/mock/mockgen/model"
)

var (
	progOnly   = flag.Bool("prog_only", false, "(reflect mode) Only generate the reflection program; write it to stdout and exit.")
	execOnly   = flag.String("exec_only", "", "(reflect mode) If set, execute this reflection program.")
	buildFlags = flag.String("build_flags", "", "(reflect mode) Additional flags for go build.")
)

func writeProgram(importPath string, symbols []string) ([]byte, error) {
	var program bytes.Buffer
	data := reflectData{
		ImportPath: importPath,
		Symbols:    symbols,
	}
	if err := reflectProgram.Execute(&program, &data); err != nil {
		return nil, err
	}
	return program.Bytes(), nil
}

// run the given program and parse the output as a model.Package.
func run(program string) (*model.Package, error) {
	f, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, err
	}

	filename := f.Name()
	defer os.Remove(filename)
	if err := f.Close(); err != nil {
		return nil, err
	}

	// Run the program.
	cmd := exec.Command(program, "-output", filename)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}

	f, err = os.Open(filename)
	if err != nil {
		return nil, err
	}

	// Process output.
	var pkg model.Package
	if err := gob.NewDecoder(f).Decode(&pkg); err != nil {
		return nil, err
	}

	if err := f.Close(); err != nil {
		return nil, err
	}

	return &pkg, nil
}

// runInDir writes the given program into the given dir, runs it there, and
// parses the output as a model.Package.
func runInDir(program []byte, dir string) (*model.Package, error) {
	// We use TempDir instead of TempFile so we can control the filename.
	tmpDir, err := ioutil.TempDir(dir, "gomock_reflect_")
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := os.RemoveAll(tmpDir); err != nil {
			log.Printf("failed to remove temp directory: %s", err)
		}
	}()
	const progSource = "prog.go"
	var progBinary = "prog.bin"
	if runtime.GOOS == "windows" {
		// Windows won't execute a program unless it has a ".exe" suffix.
		progBinary += ".exe"
	}

	if err := ioutil.WriteFile(filepath.Join(tmpDir, progSource), program, 0600); err != nil {
		return nil, err
	}

	cmdArgs := []string{}
	cmdArgs = append(cmdArgs, "build")
	if *buildFlags != "" {
		cmdArgs = append(cmdArgs, strings.Split(*buildFlags, " ")...)
	}
	cmdArgs = append(cmdArgs, "-o", progBinary, progSource)

	// Build the program.
	cmd := exec.Command("go", cmdArgs...)
	cmd.Dir = tmpDir
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	return run(filepath.Join(tmpDir, progBinary))
}

// reflectMode generates mocks via reflection on an interface.
func reflectMode(importPath string, symbols []string) (*model.Package, error) {
	// TODO: sanity check arguments

	if *execOnly != "" {
		return run(*execOnly)
	}

	program, err := writeProgram(importPath, symbols)
	if err != nil {
		return nil, err
	}

	if *progOnly {
		_, _ = os.Stdout.Write(program)
		os.Exit(0)
	}

	wd, _ := os.Getwd()

	// Try to run the program in the same directory as the input package.
	if p, err := build.Import(importPath, wd, build.FindOnly); err == nil {
		dir := p.Dir
		if p, err := runInDir(program, dir); err == nil {
			return p, nil
		}
	}

	// Since that didn't work, try to run it in the current working directory.
	if p, err := runInDir(program, wd); err == nil {
		return p, nil
	}
	// Since that didn't work, try to run it in a standard temp directory.
	return runInDir(program, "")
}

type reflectData struct {
	ImportPath string
	Symbols    []string
}

// This program reflects on an interface value, and prints the
// gob encoding of a model.Package to standard output.
// JSON doesn't work because of the model.Type interface.
var reflectProgram = template.Must(template.New("program").Parse(`
package main

import (
	"encoding/gob"
	"flag"
	"fmt"
	"os"
	"path"
	"reflect"

	"github.com/golang/mock/mockgen/model"

	pkg_ {{printf "%q" .ImportPath}}
)

var output = flag.String("output", "", "The output file name, or empty to use stdout.")

func main() {
	flag.Parse()

	its := []struct{
		sym string
		typ reflect.Type
	}{
		{{range .Symbols}}
		{ {{printf "%q" .}}, reflect.TypeOf((*pkg_.{{.}})(nil)).Elem()},
		{{end}}
	}
	pkg := &model.Package{
		// NOTE: This behaves contrary to documented behaviour if the
		// package name is not the final component of the import path.
		// The reflect package doesn't expose the package name, though.
		Name: path.Base({{printf "%q" .ImportPath}}),
	}

	for _, it := range its {
		intf, err := model.InterfaceFromInterfaceType(it.typ)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Reflection: %v\n", err)
			os.Exit(1)
		}
		intf.Name = it.sym
		pkg.Interfaces = append(pkg.Interfaces, intf)
	}

	outfile := os.Stdout
	if len(*output) != 0 {
		var err error
		outfile, err = os.Create(*output)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to open output file %q", *output)
		}
		defer func() {
			if err := outfile.Close(); err != nil {
				fmt.Fprintf(os.Stderr, "failed to close output file %q", *output)
				os.Exit(1)
			}
		}()
	}

	if err := gob.NewEncoder(outfile).Encode(pkg); err != nil {
		fmt.Fprintf(os.Stderr, "gob encode: %v\n", err)
		os.Exit(1)
	}
}
`))
<commit_msg>Reflect in Current Directory First (#390)<commit_after>// Copyright 2012 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

// This file contains the model construction by reflection.

import (
	"bytes"
	"encoding/gob"
	"flag"
	"go/build"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"text/template"

	"github.com/golang/mock/mockgen/model"
)

var (
	progOnly   = flag.Bool("prog_only", false, "(reflect mode) Only generate the reflection program; write it to stdout and exit.")
	execOnly   = flag.String("exec_only", "", "(reflect mode) If set, execute this reflection program.")
	buildFlags = flag.String("build_flags", "", "(reflect mode) Additional flags for go build.")
)

func writeProgram(importPath string, symbols []string) ([]byte, error) {
	var program bytes.Buffer
	data := reflectData{
		ImportPath: importPath,
		Symbols:    symbols,
	}
	if err := reflectProgram.Execute(&program, &data); err != nil {
		return nil, err
	}
	return program.Bytes(), nil
}

// run the given program and parse the output as a model.Package.
func run(program string) (*model.Package, error) {
	f, err := ioutil.TempFile("", "")
	if err != nil {
		return nil, err
	}

	filename := f.Name()
	defer os.Remove(filename)
	if err := f.Close(); err != nil {
		return nil, err
	}

	// Run the program.
	cmd := exec.Command(program, "-output", filename)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}

	f, err = os.Open(filename)
	if err != nil {
		return nil, err
	}

	// Process output.
	var pkg model.Package
	if err := gob.NewDecoder(f).Decode(&pkg); err != nil {
		return nil, err
	}

	if err := f.Close(); err != nil {
		return nil, err
	}

	return &pkg, nil
}

// runInDir writes the given program into the given dir, runs it there, and
// parses the output as a model.Package.
func runInDir(program []byte, dir string) (*model.Package, error) {
	// We use TempDir instead of TempFile so we can control the filename.
	tmpDir, err := ioutil.TempDir(dir, "gomock_reflect_")
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := os.RemoveAll(tmpDir); err != nil {
			log.Printf("failed to remove temp directory: %s", err)
		}
	}()
	const progSource = "prog.go"
	var progBinary = "prog.bin"
	if runtime.GOOS == "windows" {
		// Windows won't execute a program unless it has a ".exe" suffix.
		progBinary += ".exe"
	}

	if err := ioutil.WriteFile(filepath.Join(tmpDir, progSource), program, 0600); err != nil {
		return nil, err
	}

	cmdArgs := []string{}
	cmdArgs = append(cmdArgs, "build")
	if *buildFlags != "" {
		cmdArgs = append(cmdArgs, strings.Split(*buildFlags, " ")...)
	}
	cmdArgs = append(cmdArgs, "-o", progBinary, progSource)

	// Build the program.
	cmd := exec.Command("go", cmdArgs...)
	cmd.Dir = tmpDir
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return nil, err
	}

	return run(filepath.Join(tmpDir, progBinary))
}

// reflectMode generates mocks via reflection on an interface.
func reflectMode(importPath string, symbols []string) (*model.Package, error) {
	// TODO: sanity check arguments

	if *execOnly != "" {
		return run(*execOnly)
	}

	program, err := writeProgram(importPath, symbols)
	if err != nil {
		return nil, err
	}

	if *progOnly {
		_, _ = os.Stdout.Write(program)
		os.Exit(0)
	}

	wd, _ := os.Getwd()

	// Try to run the reflection program  in the current working directory.
	if p, err := runInDir(program, wd); err == nil {
		return p, nil
	}

	// Try to run the program in the same directory as the input package.
	if p, err := build.Import(importPath, wd, build.FindOnly); err == nil {
		dir := p.Dir
		if p, err := runInDir(program, dir); err == nil {
			return p, nil
		}
	}

	// Try to run it in a standard temp directory.
	return runInDir(program, "")
}

type reflectData struct {
	ImportPath string
	Symbols    []string
}

// This program reflects on an interface value, and prints the
// gob encoding of a model.Package to standard output.
// JSON doesn't work because of the model.Type interface.
var reflectProgram = template.Must(template.New("program").Parse(`
package main

import (
	"encoding/gob"
	"flag"
	"fmt"
	"os"
	"path"
	"reflect"

	"github.com/golang/mock/mockgen/model"

	pkg_ {{printf "%q" .ImportPath}}
)

var output = flag.String("output", "", "The output file name, or empty to use stdout.")

func main() {
	flag.Parse()

	its := []struct{
		sym string
		typ reflect.Type
	}{
		{{range .Symbols}}
		{ {{printf "%q" .}}, reflect.TypeOf((*pkg_.{{.}})(nil)).Elem()},
		{{end}}
	}
	pkg := &model.Package{
		// NOTE: This behaves contrary to documented behaviour if the
		// package name is not the final component of the import path.
		// The reflect package doesn't expose the package name, though.
		Name: path.Base({{printf "%q" .ImportPath}}),
	}

	for _, it := range its {
		intf, err := model.InterfaceFromInterfaceType(it.typ)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Reflection: %v\n", err)
			os.Exit(1)
		}
		intf.Name = it.sym
		pkg.Interfaces = append(pkg.Interfaces, intf)
	}

	outfile := os.Stdout
	if len(*output) != 0 {
		var err error
		outfile, err = os.Create(*output)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to open output file %q", *output)
		}
		defer func() {
			if err := outfile.Close(); err != nil {
				fmt.Fprintf(os.Stderr, "failed to close output file %q", *output)
				os.Exit(1)
			}
		}()
	}

	if err := gob.NewEncoder(outfile).Encode(pkg); err != nil {
		fmt.Fprintf(os.Stderr, "gob encode: %v\n", err)
		os.Exit(1)
	}
}
`))
<|endoftext|>


<commit_before>// Copyright 2013 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package model

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
)

const (
	// ExportedLabelPrefix is the prefix to prepend to the label names present in
	// exported metrics if a label of the same name is added by the server.
	ExportedLabelPrefix LabelName = "exported_"

	// MetricNameLabel is the label name indicating the metric name of a
	// timeseries.
	MetricNameLabel LabelName = "__name__"

	// SchemeLabel is the name of the label that holds the scheme on which to
	// scrape a target.
	SchemeLabel LabelName = "__scheme__"

	// AddressLabel is the name of the label that holds the address of
	// a scrape target.
	AddressLabel LabelName = "__address__"

	// MetricsPathLabel is the name of the label that holds the path on which to
	// scrape a target.
	MetricsPathLabel LabelName = "__metrics_path__"

	// ReservedLabelPrefix is a prefix which is not legal in user-supplied
	// label names.
	ReservedLabelPrefix = "__"

	// MetaLabelPrefix is a prefix for labels that provide meta information.
	// Labels with this prefix are used for intermediate label processing and
	// will not be attached to time series.
	MetaLabelPrefix = "__meta_"

	// ParamLabelPrefix is a prefix for labels that provide URL parameters
	// used to scrape a target.
	ParamLabelPrefix = "__param_"

	// JobLabel is the label name indicating the job from which a timeseries
	// was scraped.
	JobLabel LabelName = "job"

	// InstanceLabel is the label name used for the instance label.
	InstanceLabel LabelName = "instance"

	// BucketLabel is used for the label that defines the upper bound of a
	// bucket of a histogram ("le" -> "less or equal").
	BucketLabel = "le"

	// QuantileLabel is used for the label that defines the quantile in a
	// summary.
	QuantileLabel = "quantile"
)

// LabelNameRE is a regular expression matching valid label names.
var LabelNameRE = regexp.MustCompile("^[a-zA-Z_][a-zA-Z0-9_]*$")

// A LabelName is a key for a LabelSet or Metric.  It has a value associated
// therewith.
type LabelName string

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (ln *LabelName) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var s string
	if err := unmarshal(&s); err != nil {
		return err
	}
	if !LabelNameRE.MatchString(s) {
		return fmt.Errorf("%q is not a valid label name", s)
	}
	*ln = LabelName(s)
	return nil
}

// UnmarshalJSON implements the json.Unmarshaler interface.
func (ln *LabelName) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	if !LabelNameRE.MatchString(s) {
		return fmt.Errorf("%q is not a valid label name", s)
	}
	*ln = LabelName(s)
	return nil
}

// LabelNames is a sortable LabelName slice. In implements sort.Interface.
type LabelNames []LabelName

func (l LabelNames) Len() int {
	return len(l)
}

func (l LabelNames) Less(i, j int) bool {
	return l[i] < l[j]
}

func (l LabelNames) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func (l LabelNames) String() string {
	labelStrings := make([]string, 0, len(l))
	for _, label := range l {
		labelStrings = append(labelStrings, string(label))
	}
	return strings.Join(labelStrings, ", ")
}
<commit_msg>Reserve a label prefix for use for temporary relabeling.<commit_after>// Copyright 2013 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package model

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
)

const (
	// ExportedLabelPrefix is the prefix to prepend to the label names present in
	// exported metrics if a label of the same name is added by the server.
	ExportedLabelPrefix LabelName = "exported_"

	// MetricNameLabel is the label name indicating the metric name of a
	// timeseries.
	MetricNameLabel LabelName = "__name__"

	// SchemeLabel is the name of the label that holds the scheme on which to
	// scrape a target.
	SchemeLabel LabelName = "__scheme__"

	// AddressLabel is the name of the label that holds the address of
	// a scrape target.
	AddressLabel LabelName = "__address__"

	// MetricsPathLabel is the name of the label that holds the path on which to
	// scrape a target.
	MetricsPathLabel LabelName = "__metrics_path__"

	// ReservedLabelPrefix is a prefix which is not legal in user-supplied
	// label names.
	ReservedLabelPrefix = "__"

	// MetaLabelPrefix is a prefix for labels that provide meta information.
	// Labels with this prefix are used for intermediate label processing and
	// will not be attached to time series.
	MetaLabelPrefix = "__meta_"

	// TmpLabelPrefix is a prefix for temporary labels as part of relabelling.
	// Labels with this prefix are used for intermediate label processing and
	// will not be attached to time series. This is reserved for use in
	// Prometheus configuration files by users.
	TmpLabelPrefix = "__tmp_"

	// ParamLabelPrefix is a prefix for labels that provide URL parameters
	// used to scrape a target.
	ParamLabelPrefix = "__param_"

	// JobLabel is the label name indicating the job from which a timeseries
	// was scraped.
	JobLabel LabelName = "job"

	// InstanceLabel is the label name used for the instance label.
	InstanceLabel LabelName = "instance"

	// BucketLabel is used for the label that defines the upper bound of a
	// bucket of a histogram ("le" -> "less or equal").
	BucketLabel = "le"

	// QuantileLabel is used for the label that defines the quantile in a
	// summary.
	QuantileLabel = "quantile"
)

// LabelNameRE is a regular expression matching valid label names.
var LabelNameRE = regexp.MustCompile("^[a-zA-Z_][a-zA-Z0-9_]*$")

// A LabelName is a key for a LabelSet or Metric.  It has a value associated
// therewith.
type LabelName string

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (ln *LabelName) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var s string
	if err := unmarshal(&s); err != nil {
		return err
	}
	if !LabelNameRE.MatchString(s) {
		return fmt.Errorf("%q is not a valid label name", s)
	}
	*ln = LabelName(s)
	return nil
}

// UnmarshalJSON implements the json.Unmarshaler interface.
func (ln *LabelName) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	if !LabelNameRE.MatchString(s) {
		return fmt.Errorf("%q is not a valid label name", s)
	}
	*ln = LabelName(s)
	return nil
}

// LabelNames is a sortable LabelName slice. In implements sort.Interface.
type LabelNames []LabelName

func (l LabelNames) Len() int {
	return len(l)
}

func (l LabelNames) Less(i, j int) bool {
	return l[i] < l[j]
}

func (l LabelNames) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func (l LabelNames) String() string {
	labelStrings := make([]string, 0, len(l))
	for _, label := range l {
		labelStrings = append(labelStrings, string(label))
	}
	return strings.Join(labelStrings, ", ")
}
<|endoftext|>


<commit_before>package service

import (
	"github.com/Comcast/webpa-common/logging"
	"github.com/strava/go.serversets"
	"strings"
	"time"
)

const (
	DefaultScheme        = "http"
	DefaultHost          = "localhost"
	DefaultServer        = "localhost:2181"
	DefaultTimeout       = 5 * time.Second
	DefaultBaseDirectory = "/webpa"
	DefaultMemberPrefix  = "webpa_"
	DefaultEnvironment   = serversets.Local
	DefaultServiceName   = "test"
	DefaultVnodeCount    = 10000
)

// Options represents the set of configurable attributes for service discovery and registration
type Options struct {
	// Logger is used by any component configured via this Options.  If unset, a default
	// logger is used.
	Logger logging.Logger `json:"-"`

	// Connection is the comma-delimited Zookeeper connection string.  Both this and
	// Servers may be set, and they will be merged together when connecting to Zookeeper.
	Connection string `json:"connection,omitempty"`

	// Servers is the array of Zookeeper servers.  Both this and Connection may be set,
	// and they will be merged together when connecting to Zookeeper.
	Servers []string `json:"servers,omitempty"`

	// Timeout is the Zookeeper connection timeout.
	Timeout time.Duration `json:"timeout"`

	// BaseDirectory is the base path for all znodes created via this Options.
	BaseDirectory string `json:"baseDirectory,omitempty"`

	// MemberPrefix is the prefix for ephemeral nodes regstered via this Options.
	MemberPrefix string `json:"memberPrefix,omitempty"`

	// Environment is the environment component of the ephemeral znode path.
	Environment string `json:"environment,omitempty"`

	// ServiceName is the name of the service being registered.
	ServiceName string `json:"serviceName,omitempty"`

	// Registrations holds the slice of information used to register endpoints.  Typically,
	// this slice will either (1) be empty for an application that only watches for changes,  or (2) have the single
	// Registration indicating how this service is known.  Multiple registrations, essentially
	// being aliases for the same application, are supported.
	Registrations []string `json:"registrations,omitempty"`

	// VnodeCount is used to tune the underlying consistent hash algorithm for servers.
	VnodeCount uint `json:"vnodeCount"`

	// PingFunc is the callback function used to determine if this application is still able
	// to respond to requests.  This can be nil, and there is no default.
	PingFunc func() error `json:"-"`
}

func (o *Options) logger() logging.Logger {
	if o != nil && o.Logger != nil {
		return o.Logger
	}

	return logging.DefaultLogger()
}

func (o *Options) servers() []string {
	servers := make([]string, 0, 10)

	if o != nil {
		if len(o.Connection) > 0 {
			for _, server := range strings.Split(o.Connection, ",") {
				servers = append(servers, strings.TrimSpace(server))
			}
		}

		if len(o.Servers) > 0 {
			servers = append(servers, o.Servers...)
		}
	}

	if len(servers) == 0 {
		servers = append(servers, DefaultServer)
	}

	return servers
}

func (o *Options) timeout() time.Duration {
	if o != nil && o.Timeout > 0 {
		return time.Duration(o.Timeout)
	}

	return DefaultTimeout
}

func (o *Options) baseDirectory() string {
	if o != nil && len(o.BaseDirectory) > 0 {
		return o.BaseDirectory
	}

	return DefaultBaseDirectory
}

func (o *Options) memberPrefix() string {
	if o != nil && len(o.MemberPrefix) > 0 {
		return o.MemberPrefix
	}

	return DefaultMemberPrefix
}

func (o *Options) environment() serversets.Environment {
	if o != nil && len(o.Environment) > 0 {
		return serversets.Environment(o.Environment)
	}

	return DefaultEnvironment
}

func (o *Options) serviceName() string {
	if o != nil && len(o.ServiceName) > 0 {
		return o.ServiceName
	}

	return DefaultServiceName
}

func (o *Options) registrations() []string {
	if o != nil {
		return o.Registrations
	}

	return nil
}

func (o *Options) vnodeCount() int {
	if o != nil && o.VnodeCount > 0 {
		return int(o.VnodeCount)
	}

	return DefaultVnodeCount
}

func (o *Options) pingFunc() func() error {
	if o != nil {
		return o.PingFunc
	}

	return nil
}
<commit_msg>Changed the default vnode count to the one determined to be optimal for production<commit_after>package service

import (
	"github.com/Comcast/webpa-common/logging"
	"github.com/strava/go.serversets"
	"strings"
	"time"
)

const (
	DefaultScheme        = "http"
	DefaultHost          = "localhost"
	DefaultServer        = "localhost:2181"
	DefaultTimeout       = 5 * time.Second
	DefaultBaseDirectory = "/webpa"
	DefaultMemberPrefix  = "webpa_"
	DefaultEnvironment   = serversets.Local
	DefaultServiceName   = "test"
	DefaultVnodeCount    = 211
)

// Options represents the set of configurable attributes for service discovery and registration
type Options struct {
	// Logger is used by any component configured via this Options.  If unset, a default
	// logger is used.
	Logger logging.Logger `json:"-"`

	// Connection is the comma-delimited Zookeeper connection string.  Both this and
	// Servers may be set, and they will be merged together when connecting to Zookeeper.
	Connection string `json:"connection,omitempty"`

	// Servers is the array of Zookeeper servers.  Both this and Connection may be set,
	// and they will be merged together when connecting to Zookeeper.
	Servers []string `json:"servers,omitempty"`

	// Timeout is the Zookeeper connection timeout.
	Timeout time.Duration `json:"timeout"`

	// BaseDirectory is the base path for all znodes created via this Options.
	BaseDirectory string `json:"baseDirectory,omitempty"`

	// MemberPrefix is the prefix for ephemeral nodes regstered via this Options.
	MemberPrefix string `json:"memberPrefix,omitempty"`

	// Environment is the environment component of the ephemeral znode path.
	Environment string `json:"environment,omitempty"`

	// ServiceName is the name of the service being registered.
	ServiceName string `json:"serviceName,omitempty"`

	// Registrations holds the slice of information used to register endpoints.  Typically,
	// this slice will either (1) be empty for an application that only watches for changes,  or (2) have the single
	// Registration indicating how this service is known.  Multiple registrations, essentially
	// being aliases for the same application, are supported.
	Registrations []string `json:"registrations,omitempty"`

	// VnodeCount is used to tune the underlying consistent hash algorithm for servers.
	VnodeCount uint `json:"vnodeCount"`

	// PingFunc is the callback function used to determine if this application is still able
	// to respond to requests.  This can be nil, and there is no default.
	PingFunc func() error `json:"-"`
}

func (o *Options) logger() logging.Logger {
	if o != nil && o.Logger != nil {
		return o.Logger
	}

	return logging.DefaultLogger()
}

func (o *Options) servers() []string {
	servers := make([]string, 0, 10)

	if o != nil {
		if len(o.Connection) > 0 {
			for _, server := range strings.Split(o.Connection, ",") {
				servers = append(servers, strings.TrimSpace(server))
			}
		}

		if len(o.Servers) > 0 {
			servers = append(servers, o.Servers...)
		}
	}

	if len(servers) == 0 {
		servers = append(servers, DefaultServer)
	}

	return servers
}

func (o *Options) timeout() time.Duration {
	if o != nil && o.Timeout > 0 {
		return time.Duration(o.Timeout)
	}

	return DefaultTimeout
}

func (o *Options) baseDirectory() string {
	if o != nil && len(o.BaseDirectory) > 0 {
		return o.BaseDirectory
	}

	return DefaultBaseDirectory
}

func (o *Options) memberPrefix() string {
	if o != nil && len(o.MemberPrefix) > 0 {
		return o.MemberPrefix
	}

	return DefaultMemberPrefix
}

func (o *Options) environment() serversets.Environment {
	if o != nil && len(o.Environment) > 0 {
		return serversets.Environment(o.Environment)
	}

	return DefaultEnvironment
}

func (o *Options) serviceName() string {
	if o != nil && len(o.ServiceName) > 0 {
		return o.ServiceName
	}

	return DefaultServiceName
}

func (o *Options) registrations() []string {
	if o != nil {
		return o.Registrations
	}

	return nil
}

func (o *Options) vnodeCount() int {
	if o != nil && o.VnodeCount > 0 {
		return int(o.VnodeCount)
	}

	return DefaultVnodeCount
}

func (o *Options) pingFunc() func() error {
	if o != nil {
		return o.PingFunc
	}

	return nil
}
<|endoftext|>


<commit_before>package service

import (
	"container/heap"
	"github.com/ryszard/goskiplist/skiplist"
	"github.com/vlad-doru/fuzzyguy/levenshtein"
	"sort"
	"sync"
)

type Service interface {
	Add(key, value string)
	Get(key string) (string, bool)
	Query(key string, distance, max_results int) []string
	Len() int
}

type Storage struct {
	key      string
	value    string
	extended uint64
}

type FuzzyService struct {
	dictionary map[int]map[uint32][]Storage
	keyList    map[int]*skiplist.Set
}

func NewFuzzyService() *FuzzyService {
	dict := make(map[int]map[uint32][]Storage)
	histo := make(map[int]*skiplist.Set)
	return &FuzzyService{dict, histo}
}

func (service FuzzyService) Add(key, value string) {
	histogram := levenshtein.ComputeHistogram(key)
	storage := Storage{key, value, levenshtein.ComputeExtendedHistogram(key)}
	bucket, present := service.dictionary[len(key)]
	if present {
		list, histogram_present := bucket[histogram]
		if histogram_present {
			for _, pair := range list {
				if pair.key == key {
					pair.value = value
					return
				}
			}
			bucket[histogram] = append(bucket[histogram], storage)
			return
		}
		bucket[histogram] = []Storage{storage}
		return
	} else {
		service.keyList[len(key)] = skiplist.NewIntSet()
	}
	bucket = map[uint32][]Storage{histogram: []Storage{storage}}
	service.dictionary[len(key)] = bucket
	service.keyList[len(key)].Add(histogram)
}

func (service FuzzyService) Get(key string) (string, bool) {
	histogram := levenshtein.ComputeHistogram(key)
	bucket, present := service.dictionary[len(key)]
	if present {
		list, histogram_present := bucket[histogram]
		if histogram_present {
			for _, pair := range list {
				if pair.key == key {
					return pair.value, true
				}
			}
		}
	}
	return "", false
}

func (service FuzzyService) Len() int {
	result := 0
	for _, dict := range service.dictionary {
		result += len(dict)
	}
	return result
}

type KeyScore struct {
	score int
	key   string
}

type KeyScoreHeap []KeyScore

// We are going to implement a KeyScore max-heap based on the score
func (h KeyScoreHeap) Len() int {
	return len(h)
}

func (h KeyScoreHeap) Less(i, j int) bool {
	if h[i].score == h[j].score {
		return h[i].key < h[j].key
	}
	return h[i].score > h[j].score // this is the max-heap condition
}

func (h KeyScoreHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *KeyScoreHeap) Push(x interface{}) {
	*h = append(*h, x.(KeyScore))
}

func (h *KeyScoreHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func Abs(x int) int {
	if x < 0 {
		return -x
	}
	return x
}

func (service FuzzyService) Query(query string, threshold, max_results int) []string {
	h := new(KeyScoreHeap)
	heap.Init(h)
	query_histogram := levenshtein.ComputeHistogram(query)
	query_extended := levenshtein.ComputeExtendedHistogram(query)
	query_len := len(query)
	heap_mutex := &sync.Mutex{}
	sync_channel := make(chan int)
	start := query_len - threshold
	stop := query_len + threshold + 1

	for i := start; i < stop; i++ {
		go func(index int, mutex *sync.Mutex) {
			diff := Abs(index - query_len)
			for histogram, list := range service.dictionary[index] {
				if levenshtein.LowerBound(query_histogram, histogram, diff) > threshold {
					continue
				}
				for _, pair := range list {
					if levenshtein.ExtendedLowerBound(query_extended, pair.extended, diff) > threshold {
						continue
					}
					distance, within := levenshtein.DistanceThreshold(query, pair.key, threshold)
					if within {
						mutex.Lock()
						heap.Push(h, KeyScore{distance, pair.key})
						if h.Len() > max_results {
							heap.Pop(h)
						}
						mutex.Unlock()
					}
				}
			}
			sync_channel <- 1
		}(i, heap_mutex)
	}
	for i := start; i < stop; i++ {
		<-sync_channel
	}

	sort.Sort(h)
	results := make([]string, h.Len())
	for i := 0; i < len(results); i++ {
		results[i] = h.Pop().(KeyScore).key
	}
	return results
}
<commit_msg>Removed skiplist which was useless<commit_after>package service

import (
	"container/heap"
	"github.com/vlad-doru/fuzzyguy/levenshtein"
	"sort"
	"sync"
)

type Service interface {
	Add(key, value string)
	Get(key string) (string, bool)
	Query(key string, distance, max_results int) []string
	Len() int
}

type Storage struct {
	key      string
	value    string
	extended uint64
}

type FuzzyService struct {
	dictionary map[int]map[uint32][]Storage
}

func NewFuzzyService() *FuzzyService {
	dict := make(map[int]map[uint32][]Storage)
	return &FuzzyService{dict}
}

func (service FuzzyService) Add(key, value string) {
	histogram := levenshtein.ComputeHistogram(key)
	storage := Storage{key, value, levenshtein.ComputeExtendedHistogram(key)}
	bucket, present := service.dictionary[len(key)]
	if present {
		list, histogram_present := bucket[histogram]
		if histogram_present {
			for _, pair := range list {
				if pair.key == key {
					pair.value = value
					return
				}
			}
			bucket[histogram] = append(bucket[histogram], storage)
			return
		}
		bucket[histogram] = []Storage{storage}
		return
	}

	bucket = map[uint32][]Storage{histogram: []Storage{storage}}
	service.dictionary[len(key)] = bucket
}

func (service FuzzyService) Get(key string) (string, bool) {
	histogram := levenshtein.ComputeHistogram(key)
	bucket, present := service.dictionary[len(key)]
	if present {
		list, histogram_present := bucket[histogram]
		if histogram_present {
			for _, pair := range list {
				if pair.key == key {
					return pair.value, true
				}
			}
		}
	}
	return "", false
}

func (service FuzzyService) Len() int {
	result := 0
	for _, dict := range service.dictionary {
		result += len(dict)
	}
	return result
}

type KeyScore struct {
	score int
	key   string
}

type KeyScoreHeap []KeyScore

// We are going to implement a KeyScore max-heap based on the score
func (h KeyScoreHeap) Len() int {
	return len(h)
}

func (h KeyScoreHeap) Less(i, j int) bool {
	if h[i].score == h[j].score {
		return h[i].key < h[j].key
	}
	return h[i].score > h[j].score // this is the max-heap condition
}

func (h KeyScoreHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *KeyScoreHeap) Push(x interface{}) {
	*h = append(*h, x.(KeyScore))
}

func (h *KeyScoreHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func Abs(x int) int {
	if x < 0 {
		return -x
	}
	return x
}

func (service FuzzyService) Query(query string, threshold, max_results int) []string {
	h := new(KeyScoreHeap)
	heap.Init(h)
	query_histogram := levenshtein.ComputeHistogram(query)
	query_extended := levenshtein.ComputeExtendedHistogram(query)
	query_len := len(query)
	heap_mutex := &sync.Mutex{}
	sync_channel := make(chan int)
	start := query_len - threshold
	stop := query_len + threshold + 1

	for i := start; i < stop; i++ {
		go func(index int, mutex *sync.Mutex) {
			diff := Abs(index - query_len)
			for histogram, list := range service.dictionary[index] {
				if levenshtein.LowerBound(query_histogram, histogram, diff) > threshold {
					continue
				}
				for _, pair := range list {
					if levenshtein.ExtendedLowerBound(query_extended, pair.extended, diff) > threshold {
						continue
					}
					distance, within := levenshtein.DistanceThreshold(query, pair.key, threshold)
					if within {
						mutex.Lock()
						heap.Push(h, KeyScore{distance, pair.key})
						if h.Len() > max_results {
							heap.Pop(h)
						}
						mutex.Unlock()
					}
				}
			}
			sync_channel <- 1
		}(i, heap_mutex)
	}
	for i := start; i < stop; i++ {
		<-sync_channel
	}

	sort.Sort(h)
	results := make([]string, h.Len())
	for i := 0; i < len(results); i++ {
		results[i] = h.Pop().(KeyScore).key
	}
	return results
}
<|endoftext|>


<commit_before>package models

import (
	"github.com/techjanitor/pram-get/config"
	e "github.com/techjanitor/pram-get/errors"
	u "github.com/techjanitor/pram-get/utils"
)

// FavoriteModel holds the parameters from the request and also the key for the cache
type FavoriteModel struct {
	User   uint
	Id     uint
	Result FavoriteType
}

// IndexType is the top level of the JSON response
type FavoriteType struct {
	Starred bool `json:"starred"`
}

// Get will gather the information from the database and return it as JSON serialized data
func (i *FavoriteModel) Get() (err error) {

	// Initialize response header
	response := FavoriteType{}

	// Get Database handle
	db, err := u.GetDb()
	if err != nil {
		return
	}

	// see if a user has starred an image
	err = db.QueryRow("select count(*) from favorites where user_id = ? AND image_id = ? LIMIT 1", i.User, i.Id).Scan(&response.Starred)
	if err != nil {
		return
	}

	// This is the data we will serialize
	i.Result = response

	return

}
<commit_msg>add star check page<commit_after>package models

import (
	u "github.com/techjanitor/pram-get/utils"
)

// FavoriteModel holds the parameters from the request and also the key for the cache
type FavoriteModel struct {
	User   uint
	Id     uint
	Result FavoriteType
}

// IndexType is the top level of the JSON response
type FavoriteType struct {
	Starred bool `json:"starred"`
}

// Get will gather the information from the database and return it as JSON serialized data
func (i *FavoriteModel) Get() (err error) {

	// Initialize response header
	response := FavoriteType{}

	// Get Database handle
	db, err := u.GetDb()
	if err != nil {
		return
	}

	// see if a user has starred an image
	err = db.QueryRow("select count(*) from favorites where user_id = ? AND image_id = ? LIMIT 1", i.User, i.Id).Scan(&response.Starred)
	if err != nil {
		return
	}

	// This is the data we will serialize
	i.Result = response

	return

}
<|endoftext|>


<commit_before>package module

import (
	"archive/zip"
	"io"
	"net/http"
	"strings"

	"../utility/array"
)

type archZipRoutingModule struct {
	vpath string
	files []string
	w     http.ResponseWriter
}

func (o *archZipRoutingModule) ReturnFiles() {
	// TODO
}

func (o *archZipRoutingModule) ReturnBinary() {
	http.Error(o.w, "Not Support", http.StatusUnsupportedMediaType)
}

func (o *archZipRoutingModule) Close() {
	// Nothing to do
}

func archZipRouting(r io.ReadCloser, vpath string, w http.ResponseWriter, size int64) RoutingModule {
	http.Error(w, "Not Support", http.StatusUnsupportedMediaType)
	r.Close()
	return nil
}

func archZipRouting2(path, vpath string, w http.ResponseWriter) RoutingModule {
	// open archive
	r, err := zip.OpenReader(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}
	defer r.Close()

	// get all file path
	paths := make([]string, len(r.File))
	for i, f := range r.File {
		paths[i] = strings.Replace(f.Name, "\\", "/", -1)
	}

	// vpath is include file ?
	if vpath != "" {
		for _, f := range r.File {
			if f.Name != vpath {
				continue
			}

			conf := dispatch(vpath, w)
			if conf == nil {
				return nil
			}

			file, err := f.Open()
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return nil
			}
			return conf.routing(file, "", w, int64(f.UncompressedSize64))
		}
	}

	// vpath is include directory ?
	if vpath == "" || array.IsIncludeFunc(vpath, paths, strings.HasPrefix) {
		return &archZipRoutingModule{
			vpath: vpath,
			files: paths,
			w:     w,
		}
	}

	http.Error(w, "Not Found", http.StatusInternalServerError)
	return nil
}

var zipConf = install(&moduleConfig{
	name:     "zip",
	exts:     []string{"zip"},
	routing:  archZipRouting,
	routing2: archZipRouting2,
})
<commit_msg>fix: crash when read file into zip<commit_after>package module

import (
	"archive/zip"
	"io"
	"net/http"
	"strings"

	"../utility/array"
)

type archZipRoutingModule struct {
	vpath string
	files []string
	w     http.ResponseWriter
}

func (o *archZipRoutingModule) ReturnFiles() {
	// TODO
}

func (o *archZipRoutingModule) ReturnBinary() {
	http.Error(o.w, "Not Support", http.StatusUnsupportedMediaType)
}

func (o *archZipRoutingModule) Close() {
	// Nothing to do
}

type zipFileReadCloser struct {
	r *zip.ReadCloser
	f io.ReadCloser
}

func (o *zipFileReadCloser) Read(p []byte) (n int, err error) {
	return o.f.Read(p)
}

func (o *zipFileReadCloser) Close() error {
	o.f.Close()
	return o.r.Close()
}

func archZipRouting(r io.ReadCloser, vpath string, w http.ResponseWriter, size int64) RoutingModule {
	http.Error(w, "Not Support", http.StatusUnsupportedMediaType)
	r.Close()
	return nil
}

func archZipRouting2(path, vpath string, w http.ResponseWriter) RoutingModule {
	// open archive
	r, err := zip.OpenReader(path)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil
	}

	// get all file path
	paths := make([]string, len(r.File))
	for i, f := range r.File {
		paths[i] = strings.Replace(f.Name, "\\", "/", -1)
	}

	// vpath is include file ?
	if vpath != "" {
		for _, f := range r.File {
			if f.Name != vpath {
				continue
			}

			conf := dispatch(vpath, w)
			if conf == nil {
				r.Close()
				return nil
			}

			file, err := f.Open()
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				r.Close()
				return nil
			}

			zipfile := &zipFileReadCloser{
				r: r,
				f: file,
			}
			return conf.routing(zipfile, "", w, int64(f.UncompressedSize64))
		}
	}

	// vpath is include directory ?
	if vpath == "" || array.IsIncludeFunc(vpath, paths, strings.HasPrefix) {
		return &archZipRoutingModule{
			vpath: vpath,
			files: paths,
			w:     w,
		}
	}

	http.Error(w, "Not Found", http.StatusInternalServerError)
	return nil
}

var zipConf = install(&moduleConfig{
	name:     "zip",
	exts:     []string{"zip"},
	routing:  archZipRouting,
	routing2: archZipRouting2,
})
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"text/template"
	"time"
)

type JobStatus int

const (
	NEVER_RUN JobStatus = iota
	POLLING
	RUNNING
	FAILED
	RECOVERED
	SUCCESSFUL
)

type Job struct {
	Name       string
	Enabled    bool
	running    bool
	Git_url    string
	LastRun    time.Time
	LastStatus JobStatus
	CurStatus  JobStatus
}

func (j *Job) needsRunning() bool {
	return j.CurStatus == NEVER_RUN || j.needsUpdate()
}

func (j *Job) run() {
	log.Println("Running job:", j.Name)
	// todo: akelmore - make status a stack, not just two
	j.LastStatus = j.CurStatus
	j.CurStatus = RUNNING
	j.LastRun = time.Now()

	if j.LastStatus == NEVER_RUN {
		j.firstTimeSetup()
	} else {
		j.update()
	}

	if j.CurStatus != FAILED {
		switch j.LastStatus {
		case FAILED:
			j.CurStatus = RECOVERED
		default:
			j.CurStatus = SUCCESSFUL
		}
	}

	saveConfig()
}

func (j *Job) firstTimeSetup() {
	log.Println("Running first time setup for:", j.Name)

	var b bytes.Buffer
	multi := io.MultiWriter(&b, os.Stdout)

	cmd := exec.Command("git", "-c", "jobs/"+j.Name, "clone", "--depth", "1", j.Git_url)
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		log.Println("Error doing first time setup for:", j.Name)
		j.CurStatus = FAILED
	}
}

func (j *Job) needsUpdate() bool {
	if time.Since(j.LastRun) < 30*time.Second {
		return false
	}
	log.Println("Running needsUpdate for:", j.Name)

	// todo: akelmore - pull this multiwriter into Job so it can be output on the web
	var b bytes.Buffer
	multi := io.MultiWriter(&b, os.Stdout)

	cmd := exec.Command("git", "-c", "jobs/"+j.Name, "ls-remote", "origin", "-h", "HEAD")
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		return false
	}

	remoteHead := string(bytes.Fields(b.Bytes())[0])

	b.Reset()
	cmd = exec.Command("git", "-c", "jobs/"+j.Name, "rev-parse", "HEAD")
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		return false
	}

	localHead := string(bytes.Fields(b.Bytes())[0])

	return remoteHead != localHead
}

func (j *Job) update() {
	log.Println("Running update for:", j.Name)

	var b bytes.Buffer
	multi := io.MultiWriter(&b, os.Stdout)

	cmd := exec.Command("git", "-c", "jobs/"+j.Name, "pull", "--depth", "1")
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		log.Println("Error pulling git for:", j.Name)
		j.CurStatus = FAILED
	} else if bytes.Contains(b.Bytes(), []byte("Already up-to-date.")) {
		log.Println("Something went wrong with the git pull, it was already up to date. It shouldn't have been.")
		j.CurStatus = FAILED
	}
}

type CatarangConfig struct {
	Jobs []Job
}

var config CatarangConfig
var config_file_name = "catarang_config.json"

func addJob(w http.ResponseWriter, r *http.Request) {
	job := Job{Enabled: true}
	job.Name = r.FormValue("name")
	job.Git_url = r.FormValue("git_url")
	config.Jobs = append(config.Jobs, job)
	saveConfig()

	renderWebpage(w, r)
}

func deleteJob(w http.ResponseWriter, r *http.Request) {
	renderWebpage(w, r)
}

func pollJobs() {
	for {
		for index := range config.Jobs {
			if config.Jobs[index].needsRunning() {
				config.Jobs[index].run()
			}
		}
		time.Sleep(time.Second * 5)
	}
}

func renderWebpage(w http.ResponseWriter, r *http.Request) {
	root, err := template.ParseFiles("root.html")
	if err != nil {
		log.Println("Can't parse root.html file.")
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	root.Execute(w, config)
}

// todo: akelmore - fix threading with the reading/writing of the config
func readInConfig() {
	data, err := ioutil.ReadFile(config_file_name)
	if err != nil {
		log.Printf("Couldn't find %v, using default values.\n", config_file_name)
		return
	}

	if err = json.Unmarshal(data, &config); err != nil {
		log.Println("Error reading in", config_file_name)
		log.Println(err.Error())
	}
}

func saveConfig() {
	data, err := json.MarshalIndent(&config, "", "\t")
	if err != nil {
		log.Println("Error marshaling save data:", err.Error())
		return
	}

	err = ioutil.WriteFile(config_file_name, []byte(data), 0644)
	if err != nil {
		log.Println("Error writing config file", config_file_name)
		log.Println(err.Error())
	}
}

func main() {
	log.Println("Running Catarang!")
	readInConfig()

	go pollJobs()

	http.HandleFunc("/", renderWebpage)
	http.HandleFunc("/addjob", addJob)
	http.HandleFunc("/deletejob", deleteJob)
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<commit_msg>Added a git config, unified path, fixed clone for first path and added username/email for catarang's git. So much, yet so much more to do.<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"text/template"
	"time"
)

type GitPluginOptions struct {
	Username string
	Email    string
	Path     string
}

func (g *GitPluginOptions) path(name string) string {
	return g.Path + name + "/"
}

type JobStatus int

const (
	NEVER_RUN JobStatus = iota
	POLLING
	RUNNING
	FAILED
	RECOVERED
	SUCCESSFUL
)

type Job struct {
	Name       string
	Enabled    bool
	running    bool
	Git_url    string
	LastRun    time.Time
	LastStatus JobStatus
	CurStatus  JobStatus
}

// make git have a username and email for catarang

func (j *Job) needsRunning() bool {
	return j.CurStatus == NEVER_RUN || j.needsUpdate()
}

func (j *Job) run() {
	log.Println("Running job:", j.Name)
	// todo: akelmore - make status a stack, not just two
	j.LastStatus = j.CurStatus
	j.CurStatus = RUNNING
	j.LastRun = time.Now()

	if j.LastStatus == NEVER_RUN {
		j.firstTimeSetup()
	} else {
		j.update()
	}

	if j.CurStatus != FAILED {
		switch j.LastStatus {
		case FAILED:
			j.CurStatus = RECOVERED
		default:
			j.CurStatus = SUCCESSFUL
		}
	}

	saveConfig()
}

func (j *Job) firstTimeSetup() {
	log.Println("Running first time setup for:", j.Name)

	var b bytes.Buffer
	multi := io.MultiWriter(&b, os.Stdout)

	cmd := exec.Command("git", "clone", "--depth", "1", j.Git_url, config.Git.path(j.Name))
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		log.Println("Error doing first time setup for:", j.Name)
		j.CurStatus = FAILED
		return
	}

	b.Reset()
	cmd = exec.Command("git", "-C", config.Git.path(j.Name), "config", "user.email", config.Git.Email)
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		log.Println("Error trying to set git email for:", j.Name)
		j.CurStatus = FAILED
		// todo: akelmore - clean up
		return
	}

	b.Reset()
	cmd = exec.Command("git", "-C", config.Git.path(j.Name), "config", "user.name", config.Git.Username)
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		log.Println("Error trying to set git username for:", j.Name)
		j.CurStatus = FAILED
		// todo: akelmore - clean up
	}
}

func (j *Job) needsUpdate() bool {
	if time.Since(j.LastRun) < 30*time.Second {
		return false
	}
	log.Println("Running needsUpdate for:", j.Name)

	// todo: akelmore - pull this multiwriter into Job so it can be output on the web
	var b bytes.Buffer
	multi := io.MultiWriter(&b, os.Stdout)

	cmd := exec.Command("git", "-C", config.Git.path(j.Name), "ls-remote", "origin", "-h", "HEAD")
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		return false
	}

	remoteHead := string(bytes.Fields(b.Bytes())[0])

	b.Reset()
	cmd = exec.Command("git", "-C", config.Git.path(j.Name), "rev-parse", "HEAD")
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		return false
	}

	localHead := string(bytes.Fields(b.Bytes())[0])

	return remoteHead != localHead
}

func (j *Job) update() {
	log.Println("Running update for:", j.Name)

	var b bytes.Buffer
	multi := io.MultiWriter(&b, os.Stdout)

	cmd := exec.Command("git", "-C", config.Git.path(j.Name), "pull", "--depth=1")
	cmd.Stdout = multi
	cmd.Stderr = multi
	if err := cmd.Run(); err != nil {
		log.Println("Error pulling git for:", j.Name)
		j.CurStatus = FAILED
	} else if bytes.Contains(b.Bytes(), []byte("Already up-to-date.")) {
		log.Println("Something went wrong with the git pull, it was already up to date. It shouldn't have been.")
		j.CurStatus = FAILED
	}
}

type CatarangConfig struct {
	Jobs []Job
	Git  GitPluginOptions
}

var config CatarangConfig
var config_file_name = "catarang_config.json"

func addJob(w http.ResponseWriter, r *http.Request) {
	job := Job{Enabled: true}
	job.Name = r.FormValue("name")
	job.Git_url = r.FormValue("git_url")
	config.Jobs = append(config.Jobs, job)
	saveConfig()

	renderWebpage(w, r)
}

func deleteJob(w http.ResponseWriter, r *http.Request) {
	renderWebpage(w, r)
}

func pollJobs() {
	for {
		for index := range config.Jobs {
			if config.Jobs[index].needsRunning() {
				config.Jobs[index].run()
			}
		}
		time.Sleep(time.Second * 5)
	}
}

func renderWebpage(w http.ResponseWriter, r *http.Request) {
	root, err := template.ParseFiles("root.html")
	if err != nil {
		log.Println("Can't parse root.html file.")
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	root.Execute(w, config)
}

// todo: akelmore - fix threading with the reading/writing of the config
func readInConfig() {
	data, err := ioutil.ReadFile(config_file_name)
	if err == nil {
		if err = json.Unmarshal(data, &config); err != nil {
			log.Println("Error reading in", config_file_name)
			log.Println(err.Error())
		}
		return
	}

	// create a new config and save it out
	log.Println("No catarang config detected, creating new one.")
	config.Git.Email = "catarang@austinkelmore.com"
	config.Git.Username = "catarang"
	config.Git.Path = "jobs/"
	saveConfig()
}

func saveConfig() {
	data, err := json.MarshalIndent(&config, "", "\t")
	if err != nil {
		log.Println("Error marshaling save data:", err.Error())
		return
	}

	err = ioutil.WriteFile(config_file_name, []byte(data), 0644)
	if err != nil {
		log.Println("Error writing config file", config_file_name)
		log.Println(err.Error())
	}
}

func main() {
	log.Println("Running Catarang!")
	readInConfig()

	go pollJobs()

	http.HandleFunc("/", renderWebpage)
	http.HandleFunc("/addjob", addJob)
	http.HandleFunc("/deletejob", deleteJob)
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<|endoftext|>