<commit_before>// Copyright 2011, 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charm

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

	"launchpad.net/juju-core/utils"

	"labix.org/v2/mgo/bson"
)

// A charm URL represents charm locations such as:
//
//     cs:~joe/oneiric/wordpress
//     cs:oneiric/wordpress-42
//     local:oneiric/wordpress
//
type URL struct {
	Schema   string // "cs" or "local"
	User     string // "joe"
	Series   string // "oneiric"
	Name     string // "wordpress"
	Revision int    // -1 if unset, N otherwise
}

var (
	validUser   = regexp.MustCompile("^[a-z0-9][a-zA-Z0-9+.-]+$")
	validSeries = regexp.MustCompile("^[a-z]+([a-z-]+[a-z])?$")
	validName   = regexp.MustCompile("^[a-z][a-z0-9]*(-[a-z0-9]*[a-z][a-z0-9]*)*$")
)

// IsValidUser returns whether user is a valid username in charm URLs.
func IsValidUser(user string) bool {
	return validUser.MatchString(user)
}

// IsValidSeries returns whether series is a valid series in charm URLs.
func IsValidSeries(series string) bool {
	return validSeries.MatchString(series)
}

// IsValidName returns whether name is a valid charm name.
func IsValidName(name string) bool {
	return validName.MatchString(name)
}

// WithRevision returns a URL equivalent to url but with Revision set
// to revision.
func (url *URL) WithRevision(revision int) *URL {
	urlCopy := *url
	urlCopy.Revision = revision
	return &urlCopy
}

// MustParseURL works like ParseURL, but panics in case of errors.
func MustParseURL(url string) *URL {
	u, err := ParseURL(url)
	if err != nil {
		panic(err)
	}
	return u
}

// ParseURL parses the provided charm URL string into its respective
// structure.
func ParseURL(url string) (*URL, error) {
	u := &URL{}
	i := strings.Index(url, ":")
	if i > 0 {
		u.Schema = url[:i]
		i++
	}
	// cs: or local:
	if u.Schema != "cs" && u.Schema != "local" {
		return nil, fmt.Errorf("charm URL has invalid schema: %q", url)
	}
	parts := strings.Split(url[i:], "/")
	if len(parts) < 1 || len(parts) > 3 {
		return nil, fmt.Errorf("charm URL has invalid form: %q", url)
	}

	// ~<username>
	if strings.HasPrefix(parts[0], "~") {
		if u.Schema == "local" {
			return nil, fmt.Errorf("local charm URL with user name: %q", url)
		}
		u.User = parts[0][1:]
		if !IsValidUser(u.User) {
			return nil, fmt.Errorf("charm URL has invalid user name: %q", url)
		}
		parts = parts[1:]
	}

	// <series>
	if len(parts) < 2 {
		return nil, fmt.Errorf("charm URL without series: %q", url)
	}
	if len(parts) == 2 {
		u.Series = parts[0]
		if !IsValidSeries(u.Series) {
			return nil, fmt.Errorf("charm URL has invalid series: %q", url)
		}
		parts = parts[1:]
	}

	// <name>[-<revision>]
	u.Name = parts[0]
	u.Revision = -1
	for i := len(u.Name) - 1; i > 0; i-- {
		c := u.Name[i]
		if c >= '0' && c <= '9' {
			continue
		}
		if c == '-' && i != len(u.Name)-1 {
			var err error
			u.Revision, err = strconv.Atoi(u.Name[i+1:])
			if err != nil {
				panic(err) // We just checked it was right.
			}
			u.Name = u.Name[:i]
		}
		break
	}
	if !IsValidName(u.Name) {
		return nil, fmt.Errorf("charm URL has invalid charm name: %q", url)
	}
	return u, nil
}

// InferURL returns a charm URL inferred from src. The provided
// src may be a valid URL, in which case it is returned as-is,
// or it may be an alias in one of the following formats:
//
//    name
//    name-revision
//    series/name
//    series/name-revision
//    schema:name
//    schema:name-revision
//    cs:~user/name
//    cs:~user/name-revision
//
// The defaultSeries paramater is used to define the resulting URL
// when src does not include that information; similarly, a missing
// schema is assumed to be 'cs'.
func InferURL(src, defaultSeries string) (*URL, error) {
	if u, err := ParseURL(src); err == nil {
		// src was a valid charm URL already
		return u, nil
	}
	if strings.HasPrefix(src, "~") {
		return nil, fmt.Errorf("cannot infer charm URL with user but no schema: %q", src)
	}
	orig := src
	schema := "cs"
	if i := strings.Index(src, ":"); i != -1 {
		schema, src = src[:i], src[i+1:]
	}
	var full string
	switch parts := strings.Split(src, "/"); len(parts) {
	case 1:
		if defaultSeries == "" {
			return nil, fmt.Errorf("cannot infer charm URL for %q: no series provided", orig)
		}
		full = fmt.Sprintf("%s:%s/%s", schema, defaultSeries, src)
	case 2:
		if strings.HasPrefix(parts[0], "~") {
			if defaultSeries == "" {
				return nil, fmt.Errorf("cannot infer charm URL for %q: no series provided", orig)
			}
			full = fmt.Sprintf("%s:%s/%s/%s", schema, parts[0], defaultSeries, parts[1])
		} else {
			full = fmt.Sprintf("%s:%s", schema, src)
		}
	default:
		full = fmt.Sprintf("%s:%s", schema, src)
	}
	u, err := ParseURL(full)
	if err != nil && orig != full {
		err = fmt.Errorf("%s (URL inferred from %q)", err, orig)
	}
	return u, err
}

func (u *URL) Path() string {
	if u.User != "" {
		if u.Revision >= 0 {
			return fmt.Sprintf("~%s/%s/%s-%d", u.User, u.Series, u.Name, u.Revision)
		}
		return fmt.Sprintf("~%s/%s/%s", u.User, u.Series, u.Name)
	}
	if u.Revision >= 0 {
		return fmt.Sprintf("%s/%s-%d", u.Series, u.Name, u.Revision)
	}
	return fmt.Sprintf("%s/%s", u.Series, u.Name)
}

func (u *URL) String() string {
	return fmt.Sprintf("%s:%s", u.Schema, u.Path())
}

// GetBSON turns u into a bson.Getter so it can be saved directly
// on a MongoDB database with mgo.
func (u *URL) GetBSON() (interface{}, error) {
	if u == nil {
		return nil, nil
	}
	return u.String(), nil
}

// SetBSON turns u into a bson.Setter so it can be loaded directly
// from a MongoDB database with mgo.
func (u *URL) SetBSON(raw bson.Raw) error {
	if raw.Kind == 10 {
		return bson.SetZero
	}
	var s string
	err := raw.Unmarshal(&s)
	if err != nil {
		return err
	}
	url, err := ParseURL(s)
	if err != nil {
		return err
	}
	*u = *url
	return nil
}

var jsonNull = []byte("null")

func (u *URL) MarshalJSON() ([]byte, error) {
	if u == nil {
		panic("cannot marshal nil *charm.URL")
	}
	return json.Marshal(u.String())
}

func (u *URL) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	url, err := ParseURL(s)
	if err != nil {
		return err
	}
	*u = *url
	return nil
}

// ArchiveName returns a string that is suitable as a file name in a
// storage URL. It is constructed from the charm name, revision and a
// random UUID string.
func (u *URL) ArchiveName() (string, error) {
	uuid, err := utils.NewUUID()
	if err != nil {
		return "", err
	}
	return Quote(fmt.Sprintf("%s-%d-%s", u.Name, u.Revision, uuid)), nil
}

// Quote translates a charm url string into one which can be safely used
// in a file path.  ASCII letters, ASCII digits, dot and dash stay the
// same; other characters are translated to their hex representation
// surrounded by underscores.
func Quote(unsafe string) string {
	safe := make([]byte, 0, len(unsafe)*4)
	for i := 0; i < len(unsafe); i++ {
		b := unsafe[i]
		switch {
		case b >= 'a' && b <= 'z',
			b >= 'A' && b <= 'Z',
			b >= '0' && b <= '9',
			b == '.',
			b == '-':
			safe = append(safe, b)
		default:
			safe = append(safe, fmt.Sprintf("_%02x_", b)...)
		}
	}
	return string(safe)
}
<commit_msg>Fixed imports ordering<commit_after>// Copyright 2011, 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charm

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

	"labix.org/v2/mgo/bson"

	"launchpad.net/juju-core/utils"
)

// A charm URL represents charm locations such as:
//
//     cs:~joe/oneiric/wordpress
//     cs:oneiric/wordpress-42
//     local:oneiric/wordpress
//
type URL struct {
	Schema   string // "cs" or "local"
	User     string // "joe"
	Series   string // "oneiric"
	Name     string // "wordpress"
	Revision int    // -1 if unset, N otherwise
}

var (
	validUser   = regexp.MustCompile("^[a-z0-9][a-zA-Z0-9+.-]+$")
	validSeries = regexp.MustCompile("^[a-z]+([a-z-]+[a-z])?$")
	validName   = regexp.MustCompile("^[a-z][a-z0-9]*(-[a-z0-9]*[a-z][a-z0-9]*)*$")
)

// IsValidUser returns whether user is a valid username in charm URLs.
func IsValidUser(user string) bool {
	return validUser.MatchString(user)
}

// IsValidSeries returns whether series is a valid series in charm URLs.
func IsValidSeries(series string) bool {
	return validSeries.MatchString(series)
}

// IsValidName returns whether name is a valid charm name.
func IsValidName(name string) bool {
	return validName.MatchString(name)
}

// WithRevision returns a URL equivalent to url but with Revision set
// to revision.
func (url *URL) WithRevision(revision int) *URL {
	urlCopy := *url
	urlCopy.Revision = revision
	return &urlCopy
}

// MustParseURL works like ParseURL, but panics in case of errors.
func MustParseURL(url string) *URL {
	u, err := ParseURL(url)
	if err != nil {
		panic(err)
	}
	return u
}

// ParseURL parses the provided charm URL string into its respective
// structure.
func ParseURL(url string) (*URL, error) {
	u := &URL{}
	i := strings.Index(url, ":")
	if i > 0 {
		u.Schema = url[:i]
		i++
	}
	// cs: or local:
	if u.Schema != "cs" && u.Schema != "local" {
		return nil, fmt.Errorf("charm URL has invalid schema: %q", url)
	}
	parts := strings.Split(url[i:], "/")
	if len(parts) < 1 || len(parts) > 3 {
		return nil, fmt.Errorf("charm URL has invalid form: %q", url)
	}

	// ~<username>
	if strings.HasPrefix(parts[0], "~") {
		if u.Schema == "local" {
			return nil, fmt.Errorf("local charm URL with user name: %q", url)
		}
		u.User = parts[0][1:]
		if !IsValidUser(u.User) {
			return nil, fmt.Errorf("charm URL has invalid user name: %q", url)
		}
		parts = parts[1:]
	}

	// <series>
	if len(parts) < 2 {
		return nil, fmt.Errorf("charm URL without series: %q", url)
	}
	if len(parts) == 2 {
		u.Series = parts[0]
		if !IsValidSeries(u.Series) {
			return nil, fmt.Errorf("charm URL has invalid series: %q", url)
		}
		parts = parts[1:]
	}

	// <name>[-<revision>]
	u.Name = parts[0]
	u.Revision = -1
	for i := len(u.Name) - 1; i > 0; i-- {
		c := u.Name[i]
		if c >= '0' && c <= '9' {
			continue
		}
		if c == '-' && i != len(u.Name)-1 {
			var err error
			u.Revision, err = strconv.Atoi(u.Name[i+1:])
			if err != nil {
				panic(err) // We just checked it was right.
			}
			u.Name = u.Name[:i]
		}
		break
	}
	if !IsValidName(u.Name) {
		return nil, fmt.Errorf("charm URL has invalid charm name: %q", url)
	}
	return u, nil
}

// InferURL returns a charm URL inferred from src. The provided
// src may be a valid URL, in which case it is returned as-is,
// or it may be an alias in one of the following formats:
//
//    name
//    name-revision
//    series/name
//    series/name-revision
//    schema:name
//    schema:name-revision
//    cs:~user/name
//    cs:~user/name-revision
//
// The defaultSeries paramater is used to define the resulting URL
// when src does not include that information; similarly, a missing
// schema is assumed to be 'cs'.
func InferURL(src, defaultSeries string) (*URL, error) {
	if u, err := ParseURL(src); err == nil {
		// src was a valid charm URL already
		return u, nil
	}
	if strings.HasPrefix(src, "~") {
		return nil, fmt.Errorf("cannot infer charm URL with user but no schema: %q", src)
	}
	orig := src
	schema := "cs"
	if i := strings.Index(src, ":"); i != -1 {
		schema, src = src[:i], src[i+1:]
	}
	var full string
	switch parts := strings.Split(src, "/"); len(parts) {
	case 1:
		if defaultSeries == "" {
			return nil, fmt.Errorf("cannot infer charm URL for %q: no series provided", orig)
		}
		full = fmt.Sprintf("%s:%s/%s", schema, defaultSeries, src)
	case 2:
		if strings.HasPrefix(parts[0], "~") {
			if defaultSeries == "" {
				return nil, fmt.Errorf("cannot infer charm URL for %q: no series provided", orig)
			}
			full = fmt.Sprintf("%s:%s/%s/%s", schema, parts[0], defaultSeries, parts[1])
		} else {
			full = fmt.Sprintf("%s:%s", schema, src)
		}
	default:
		full = fmt.Sprintf("%s:%s", schema, src)
	}
	u, err := ParseURL(full)
	if err != nil && orig != full {
		err = fmt.Errorf("%s (URL inferred from %q)", err, orig)
	}
	return u, err
}

func (u *URL) Path() string {
	if u.User != "" {
		if u.Revision >= 0 {
			return fmt.Sprintf("~%s/%s/%s-%d", u.User, u.Series, u.Name, u.Revision)
		}
		return fmt.Sprintf("~%s/%s/%s", u.User, u.Series, u.Name)
	}
	if u.Revision >= 0 {
		return fmt.Sprintf("%s/%s-%d", u.Series, u.Name, u.Revision)
	}
	return fmt.Sprintf("%s/%s", u.Series, u.Name)
}

func (u *URL) String() string {
	return fmt.Sprintf("%s:%s", u.Schema, u.Path())
}

// GetBSON turns u into a bson.Getter so it can be saved directly
// on a MongoDB database with mgo.
func (u *URL) GetBSON() (interface{}, error) {
	if u == nil {
		return nil, nil
	}
	return u.String(), nil
}

// SetBSON turns u into a bson.Setter so it can be loaded directly
// from a MongoDB database with mgo.
func (u *URL) SetBSON(raw bson.Raw) error {
	if raw.Kind == 10 {
		return bson.SetZero
	}
	var s string
	err := raw.Unmarshal(&s)
	if err != nil {
		return err
	}
	url, err := ParseURL(s)
	if err != nil {
		return err
	}
	*u = *url
	return nil
}

var jsonNull = []byte("null")

func (u *URL) MarshalJSON() ([]byte, error) {
	if u == nil {
		panic("cannot marshal nil *charm.URL")
	}
	return json.Marshal(u.String())
}

func (u *URL) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	url, err := ParseURL(s)
	if err != nil {
		return err
	}
	*u = *url
	return nil
}

// ArchiveName returns a string that is suitable as a file name in a
// storage URL. It is constructed from the charm name, revision and a
// random UUID string.
func (u *URL) ArchiveName() (string, error) {
	uuid, err := utils.NewUUID()
	if err != nil {
		return "", err
	}
	return Quote(fmt.Sprintf("%s-%d-%s", u.Name, u.Revision, uuid)), nil
}

// Quote translates a charm url string into one which can be safely used
// in a file path.  ASCII letters, ASCII digits, dot and dash stay the
// same; other characters are translated to their hex representation
// surrounded by underscores.
func Quote(unsafe string) string {
	safe := make([]byte, 0, len(unsafe)*4)
	for i := 0; i < len(unsafe); i++ {
		b := unsafe[i]
		switch {
		case b >= 'a' && b <= 'z',
			b >= 'A' && b <= 'Z',
			b >= '0' && b <= '9',
			b == '.',
			b == '-':
			safe = append(safe, b)
		default:
			safe = append(safe, fmt.Sprintf("_%02x_", b)...)
		}
	}
	return string(safe)
}
<|endoftext|>


<commit_before>// Copyright 2018 The Kubeflow 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 commands

import (
	"fmt"
	"os"
	"path"
	"time"

	podlogs "github.com/kubeflow/arena/pkg/podlogs"
	tlogs "github.com/kubeflow/arena/pkg/printer/base/logs"
	"github.com/kubeflow/arena/pkg/util"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"k8s.io/client-go/kubernetes"
)

func NewLogsCommand() *cobra.Command {
	var outerArgs = &podlogs.OuterRequestArgs{}
	var command = &cobra.Command{
		Use:   "logs training job",
		Short: "print the logs for a task of the training job",
		Run: func(cmd *cobra.Command, args []string) {
			util.SetLogLevel(logLevel)
			if len(args) == 0 {
				cmd.HelpFunc()(cmd, args)
				os.Exit(1)
			}
			name = args[0]
			setupKubeconfig()
			conf, err := clientConfig.ClientConfig()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			_, err = initKubeClient()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			outerArgs.KubeClient = kubernetes.NewForConfigOrDie(conf)
			err = updateNamespace(cmd)
			if err != nil {
				log.Debugf("Failed due to %v", err)
				fmt.Println(err)
				os.Exit(1)
			}

			// podName, err := getPodNameFromJob(printer.kubeClient, namespace, name)
			job, err := searchTrainingJob(name, trainingType, namespace)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			outerArgs.Namespace = namespace
			outerArgs.RetryCount = 5
			outerArgs.RetryTimeout = time.Millisecond
			names := []string{}
			for _, pod := range job.AllPods() {
				names = append(names, path.Base(pod.ObjectMeta.SelfLink))
			}
			logPrinter, err := tlogs.NewPodLogPrinter(names, outerArgs)
			if err != nil {
				log.Errorf(err.Error())
				os.Exit(1)
			}
			code, err := logPrinter.Print()
			if err != nil {
				log.Errorf("%s, %s", err.Error(), "please use \"arena get\" to get more information.")
				os.Exit(1)
			} else if code != 0 {
				os.Exit(code)
			}
		},
	}

	command.Flags().StringVar(&trainingType, "type", "", "The training type to show logging, the possible option is tfjob, mpijob, horovodjob or standalonejob. (optional)")

	command.Flags().BoolVarP(&outerArgs.Follow, "follow", "f", false, "Specify if the logs should be streamed.")
	command.Flags().StringVar(&outerArgs.SinceSeconds, "since", "", "Only return logs newer than a relative duration like 5s, 2m, or 3h. Defaults to all logs. Only one of since-time / since may be used.")
	command.Flags().StringVar(&outerArgs.SinceTime, "since-time", "", "Only return logs after a specific date (RFC3339). Defaults to all logs. Only one of since-time / since may be used.")
	command.Flags().IntVarP(&outerArgs.Tail, "tail", "t", -1, "Lines of recent log file to display. Defaults to -1 with no selector, showing all log lines otherwise 10, if a selector is provided.")
	command.Flags().BoolVar(&outerArgs.Timestamps, "timestamps", false, "Include timestamps on each line in the log output")

	// command.Flags().StringVar(&printer.pod, "instance", "", "Only return logs after a specific date (RFC3339). Defaults to all logs. Only one of since-time / since may be used.")
	command.Flags().StringVarP(&outerArgs.PodName, "instance", "i", "", "Specify the task instance to get log")
	return command
}
<commit_msg>change the arena logs behavior (#258)<commit_after>// Copyright 2018 The Kubeflow 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 commands

import (
	"fmt"
	"os"
	"path"
	"time"

	podlogs "github.com/kubeflow/arena/pkg/podlogs"
	tlogs "github.com/kubeflow/arena/pkg/printer/base/logs"
	"github.com/kubeflow/arena/pkg/util"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"k8s.io/client-go/kubernetes"
)

func NewLogsCommand() *cobra.Command {
	var outerArgs = &podlogs.OuterRequestArgs{}
	var command = &cobra.Command{
		Use:   "logs training job",
		Short: "print the logs for a task of the training job",
		Run: func(cmd *cobra.Command, args []string) {
			util.SetLogLevel(logLevel)
			if len(args) == 0 {
				cmd.HelpFunc()(cmd, args)
				os.Exit(1)
			}
			name = args[0]
			setupKubeconfig()
			conf, err := clientConfig.ClientConfig()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			_, err = initKubeClient()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			outerArgs.KubeClient = kubernetes.NewForConfigOrDie(conf)
			err = updateNamespace(cmd)
			if err != nil {
				log.Debugf("Failed due to %v", err)
				fmt.Println(err)
				os.Exit(1)
			}

			// podName, err := getPodNameFromJob(printer.kubeClient, namespace, name)
			job, err := searchTrainingJob(name, trainingType, namespace)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
			outerArgs.Namespace = namespace
			outerArgs.RetryCount = 5
			outerArgs.RetryTimeout = time.Millisecond
			names := []string{}
			for _, pod := range job.AllPods() {
				names = append(names, path.Base(pod.ObjectMeta.SelfLink))
			}
			chiefPod := job.ChiefPod()
			if len(names) > 1  && outerArgs.PodName == "" {
				names = []string{path.Base(chiefPod.ObjectMeta.SelfLink)}
			}
			logPrinter, err := tlogs.NewPodLogPrinter(names, outerArgs)
			if err != nil {
				log.Errorf(err.Error())
				os.Exit(1)
			}
			code, err := logPrinter.Print()
			if err != nil {
				log.Errorf("%s, %s", err.Error(), "please use \"arena get\" to get more information.")
				os.Exit(1)
			} else if code != 0 {
				os.Exit(code)
			}
		},
	}

	command.Flags().StringVar(&trainingType, "type", "", "The training type to show logging, the possible option is tfjob, mpijob, horovodjob or standalonejob. (optional)")

	command.Flags().BoolVarP(&outerArgs.Follow, "follow", "f", false, "Specify if the logs should be streamed.")
	command.Flags().StringVar(&outerArgs.SinceSeconds, "since", "", "Only return logs newer than a relative duration like 5s, 2m, or 3h. Defaults to all logs. Only one of since-time / since may be used.")
	command.Flags().StringVar(&outerArgs.SinceTime, "since-time", "", "Only return logs after a specific date (RFC3339). Defaults to all logs. Only one of since-time / since may be used.")
	command.Flags().IntVarP(&outerArgs.Tail, "tail", "t", -1, "Lines of recent log file to display. Defaults to -1 with no selector, showing all log lines otherwise 10, if a selector is provided.")
	command.Flags().BoolVar(&outerArgs.Timestamps, "timestamps", false, "Include timestamps on each line in the log output")

	// command.Flags().StringVar(&printer.pod, "instance", "", "Only return logs after a specific date (RFC3339). Defaults to all logs. Only one of since-time / since may be used.")
	command.Flags().StringVarP(&outerArgs.PodName, "instance", "i", "", "Specify the task instance to get log")
	return command
}
<|endoftext|>


<commit_before>// Copyright © 2018 Banzai Cloud
//
// 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

import (
	"fmt"

	"github.com/spf13/cobra"
)

const appName = "bank-vaults"
const version = "v0.0.1"

func init() {
	rootCmd.AddCommand(versionCmd)
}

var versionCmd = &cobra.Command{
	Use:   "version",
	Short: fmt.Sprint("Print the version number of ", appName),
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Printf("%s %s\n", appName, version)
	},
}
<commit_msg>remove unused version command<commit_after><|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"

	"github.com/2opremio/go-k8s-portforward"
	"github.com/pkg/errors"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// Attempt to create PortForwards to fluxes that match the label selectors until a Flux
// is found or an error is returned.
func tryPortforwards(ns string, selectors ...metav1.LabelSelector) (p *portforward.PortForward, err error) {
	message := fmt.Sprintf("Flux pod not found for labels in namespace %s:", ns)

	for _, selector := range selectors {
		p, err = tryPortforward(ns, selector)
		if err == nil {
			return
		}

		if !strings.Contains(err.Error(), "Could not find pod for selector") {
			return
		} else {
			message = fmt.Sprintf("%s\n  %s", message, metav1.FormatLabelSelector(&selector))
		}
	}

	if err != nil {
		err = errors.New(message)
	}

	return
}

// Attempt to create a portforward in the namespace for the provided LabelSelector
func tryPortforward(ns string, selector metav1.LabelSelector) (*portforward.PortForward, error) {
	portforwarder, err := portforward.NewPortForwarder(ns, selector, 3030)
	if err != nil {
		return portforwarder, err
	}

	err = portforwarder.Start()
	if err != nil {
		return portforwarder, err
	}

	return portforwarder, nil
}
<commit_msg>Improve fluxctl connection error<commit_after>package main

import (
	"fmt"
	"strings"

	"github.com/2opremio/go-k8s-portforward"
	"github.com/pkg/errors"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// Attempt to create PortForwards to fluxes that match the label selectors until a Flux
// is found or an error is returned.
func tryPortforwards(ns string, selectors ...metav1.LabelSelector) (p *portforward.PortForward, err error) {
	message := fmt.Sprintf("No pod found in namespace %q using the following selectors:", ns)

	for _, selector := range selectors {
		p, err = tryPortforward(ns, selector)
		if err == nil {
			return
		}

		if !strings.Contains(err.Error(), "Could not find pod for selector") {
			return
		} else {
			message = fmt.Sprintf("%s\n  %s", message, metav1.FormatLabelSelector(&selector))
		}
	}
	message = fmt.Sprintf("%s\n\nMake sure Flux is running in namespace %q.\n"+
		"If Flux is running in another different namespace, please supply it to --k8s-fwd-ns.", message, ns)
	if err != nil {
		err = errors.New(message)
	}

	return
}

// Attempt to create a portforward in the namespace for the provided LabelSelector
func tryPortforward(ns string, selector metav1.LabelSelector) (*portforward.PortForward, error) {
	portforwarder, err := portforward.NewPortForwarder(ns, selector, 3030)
	if err != nil {
		return portforwarder, err
	}

	err = portforwarder.Start()
	if err != nil {
		return portforwarder, err
	}

	return portforwarder, nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"io"
	"time"

	"github.com/spf13/cobra"

	"github.com/weaveworks/flux"
	"github.com/weaveworks/flux/api/v11"
	"github.com/weaveworks/flux/api/v6"
	"github.com/weaveworks/flux/job"
	"github.com/weaveworks/flux/update"
)

type controllerReleaseOpts struct {
	*rootOpts
	namespace      string
	controllers    []string
	allControllers bool
	image          string
	allImages      bool
	exclude        []string
	dryRun         bool
	interactive    bool
	force          bool
	watch          bool
	outputOpts
	cause update.Cause

	// Deprecated
	services []string
}

func newControllerRelease(parent *rootOpts) *controllerReleaseOpts {
	return &controllerReleaseOpts{rootOpts: parent}
}

func (opts *controllerReleaseOpts) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "release",
		Short: "Release a new version of a controller.",
		Example: makeExample(
			"fluxctl release -n default --controller=deployment/foo --update-image=library/hello:v2",
			"fluxctl release --all --update-image=library/hello:v2",
			"fluxctl release --controller=default:deployment/foo --update-all-images",
		),
		RunE: opts.RunE,
	}

	AddOutputFlags(cmd, &opts.outputOpts)
	AddCauseFlags(cmd, &opts.cause)
	cmd.Flags().StringVarP(&opts.namespace, "namespace", "n", "default", "Controller namespace")
	cmd.Flags().StringSliceVarP(&opts.controllers, "controller", "c", []string{}, "List of controllers to release <namespace>:<kind>/<name>")
	cmd.Flags().BoolVar(&opts.allControllers, "all", false, "Release all controllers")
	cmd.Flags().StringVarP(&opts.image, "update-image", "i", "", "Update a specific image")
	cmd.Flags().BoolVar(&opts.allImages, "update-all-images", false, "Update all images to latest versions")
	cmd.Flags().StringSliceVar(&opts.exclude, "exclude", []string{}, "List of controllers to exclude")
	cmd.Flags().BoolVar(&opts.dryRun, "dry-run", false, "Do not release anything; just report back what would have been done")
	cmd.Flags().BoolVar(&opts.interactive, "interactive", false, "Select interactively which containers to update")
	cmd.Flags().BoolVarP(&opts.force, "force", "f", false, "Disregard locks and container image filters (has no effect when used with --all or --update-all-images)")
	cmd.Flags().BoolVarP(&opts.watch, "watch", "w", false, "Watch rollout progress during release")

	// Deprecated
	cmd.Flags().StringSliceVarP(&opts.services, "service", "s", []string{}, "Service to release")
	cmd.Flags().MarkHidden("service")

	return cmd
}

func (opts *controllerReleaseOpts) RunE(cmd *cobra.Command, args []string) error {
	if len(opts.services) > 0 {
		return errorServiceFlagDeprecated
	}

	if len(args) != 0 {
		return errorWantedNoArgs
	}

	if err := checkExactlyOne("--update-image=<image> or --update-all-images", opts.image != "", opts.allImages); err != nil {
		return err
	}

	switch {
	case len(opts.controllers) <= 0 && !opts.allControllers:
		return newUsageError("please supply either --all, or at least one --controller=<controller>")
	case opts.force && opts.allControllers && opts.allImages:
		return newUsageError("--force has no effect when used with --all and --update-all-images")
	case opts.force && opts.allControllers:
		fmt.Fprintf(cmd.OutOrStderr(), "Warning: --force will not ignore locked controllers when used with --all\n")
	case opts.force && opts.allImages:
		fmt.Fprintf(cmd.OutOrStderr(), "Warning: --force will not ignore container image tags when used with --update-all-images\n")
	}

	var (
		controllers   []update.ResourceSpec
		controllerIds []flux.ResourceID
	)
	if opts.allControllers {
		controllers = []update.ResourceSpec{update.ResourceSpecAll}
	} else {
		for _, controller := range opts.controllers {
			id, err := flux.ParseResourceIDOptionalNamespace(opts.namespace, controller)
			if err != nil {
				return err
			}
			controllerIds = append(controllerIds, id)
			controllers = append(controllers, update.MakeResourceSpec(id))
		}
	}

	var (
		image update.ImageSpec
		err   error
	)
	switch {
	case opts.image != "":
		image, err = update.ParseImageSpec(opts.image)
		if err != nil {
			return err
		}
	case opts.allImages:
		image = update.ImageSpecLatest
	}

	var kind update.ReleaseKind = update.ReleaseKindExecute
	if opts.dryRun || opts.interactive {
		kind = update.ReleaseKindPlan
	}

	var excludes []flux.ResourceID
	for _, exclude := range opts.exclude {
		s, err := flux.ParseResourceIDOptionalNamespace(opts.namespace, exclude)
		if err != nil {
			return err
		}
		excludes = append(excludes, s)
	}

	if kind == update.ReleaseKindPlan {
		fmt.Fprintf(cmd.OutOrStderr(), "Submitting dry-run release ...\n")
	} else {
		fmt.Fprintf(cmd.OutOrStderr(), "Submitting release ...\n")
	}

	ctx := context.Background()
	spec := update.ReleaseImageSpec{
		ServiceSpecs: controllers,
		ImageSpec:    image,
		Kind:         kind,
		Excludes:     excludes,
		Force:        opts.force,
	}
	jobID, err := opts.API.UpdateManifests(ctx, update.Spec{
		Type:  update.Images,
		Cause: opts.cause,
		Spec:  spec,
	})
	if err != nil {
		return err
	}

	if opts.interactive {
		result, err := awaitJob(ctx, opts.API, jobID)
		if err != nil {
			return err
		}

		spec, err := promptSpec(cmd.OutOrStdout(), result, opts.verbosity)
		spec.Force = opts.force
		if err != nil {
			fmt.Fprintln(cmd.OutOrStderr(), err.Error())
			return nil
		}

		fmt.Fprintf(cmd.OutOrStderr(), "Submitting selected release ...\n")
		jobID, err = opts.API.UpdateManifests(ctx, update.Spec{
			Type:  update.Containers,
			Cause: opts.cause,
			Spec:  spec,
		})
		if err != nil {
			fmt.Fprintln(cmd.OutOrStderr(), err.Error())
			return nil
		}

		opts.dryRun = false
	}

	err = await(ctx, cmd.OutOrStdout(), cmd.OutOrStderr(), opts.API, jobID, !opts.dryRun, opts.verbosity)
	if err != nil {
		return err
	}

	if !opts.watch {
		return nil
	}

	fmt.Fprintf(cmd.OutOrStderr(), "Monitoring rollout ...\n")
	for {
		completed := 0
		services, err := opts.API.ListServicesWithOptions(ctx, v11.ListServicesOptions{Services: controllerIds})
		if err != nil {
			return err
		}

		for _, service := range services {
			writeRolloutStatus(service, opts.verbosity)

			if service.Rollout.Outdated == 0 {
				completed++
			}

			if service.Rollout.Messages != nil {
				fmt.Fprintf(cmd.OutOrStderr(), "There was a problem releasing %s:\n", service.ID)
				for _, msg := range service.Rollout.Messages {
					fmt.Fprintf(cmd.OutOrStderr(), "%s\n", msg)
				}
				return nil
			}
		}

		if completed == len(services) {
			fmt.Fprintf(cmd.OutOrStderr(), "All controllers ready.\n")
			return nil
		}

		time.Sleep(2000 * time.Millisecond)
	}
}

func writeRolloutStatus(service v6.ControllerStatus, verbosity int) {
	w := newTabwriter()
	fmt.Fprintf(w, "CONTROLLER\tCONTAINER\tIMAGE\tRELEASE\tREPLICAS\n")

	if len(service.Containers) > 0 {
		c := service.Containers[0]
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%d/%d", service.ID, c.Name, c.Current.ID, service.Status, service.Rollout.Desired, service.Rollout.Ready)
		if verbosity > 0 {
			fmt.Fprintf(w, " (%d outdated, %d updated)", service.Rollout.Outdated, service.Rollout.Updated)
		}
		fmt.Fprintf(w, "\n")
		for _, c := range service.Containers[1:] {
			fmt.Fprintf(w, "\t%s\t%s\t\t\n", c.Name, c.Current.ID)
		}
	} else {
		fmt.Fprintf(w, "%s\t\t\t%s\t%d/%d", service.ID, service.Status, service.Rollout.Desired, service.Rollout.Ready)
		if verbosity > 0 {
			fmt.Fprintf(w, " (%d outdated, %d updated)", service.Rollout.Outdated, service.Rollout.Updated)
		}
		fmt.Fprintf(w, "\n")
	}
	fmt.Fprintln(w)
	w.Flush()
}

func promptSpec(out io.Writer, result job.Result, verbosity int) (update.ReleaseContainersSpec, error) {
	menu := update.NewMenu(out, result.Result, verbosity)
	containerSpecs, err := menu.Run()
	return update.ReleaseContainersSpec{
		Kind:           update.ReleaseKindExecute,
		ContainerSpecs: containerSpecs,
		SkipMismatches: false,
	}, err
}
<commit_msg>PR Feedback and Fixes<commit_after>package main

import (
	"context"
	"fmt"
	"io"
	"time"

	"github.com/spf13/cobra"

	"github.com/weaveworks/flux"
	"github.com/weaveworks/flux/api/v11"
	"github.com/weaveworks/flux/api/v6"
	"github.com/weaveworks/flux/cluster"
	"github.com/weaveworks/flux/job"
	"github.com/weaveworks/flux/update"
)

type controllerReleaseOpts struct {
	*rootOpts
	namespace      string
	controllers    []string
	allControllers bool
	image          string
	allImages      bool
	exclude        []string
	dryRun         bool
	interactive    bool
	force          bool
	watch          bool
	outputOpts
	cause update.Cause

	// Deprecated
	services []string
}

func newControllerRelease(parent *rootOpts) *controllerReleaseOpts {
	return &controllerReleaseOpts{rootOpts: parent}
}

func (opts *controllerReleaseOpts) Command() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "release",
		Short: "Release a new version of a controller.",
		Example: makeExample(
			"fluxctl release -n default --controller=deployment/foo --update-image=library/hello:v2",
			"fluxctl release --all --update-image=library/hello:v2",
			"fluxctl release --controller=default:deployment/foo --update-all-images",
		),
		RunE: opts.RunE,
	}

	AddOutputFlags(cmd, &opts.outputOpts)
	AddCauseFlags(cmd, &opts.cause)
	cmd.Flags().StringVarP(&opts.namespace, "namespace", "n", "default", "Controller namespace")
	cmd.Flags().StringSliceVarP(&opts.controllers, "controller", "c", []string{}, "List of controllers to release <namespace>:<kind>/<name>")
	cmd.Flags().BoolVar(&opts.allControllers, "all", false, "Release all controllers")
	cmd.Flags().StringVarP(&opts.image, "update-image", "i", "", "Update a specific image")
	cmd.Flags().BoolVar(&opts.allImages, "update-all-images", false, "Update all images to latest versions")
	cmd.Flags().StringSliceVar(&opts.exclude, "exclude", []string{}, "List of controllers to exclude")
	cmd.Flags().BoolVar(&opts.dryRun, "dry-run", false, "Do not release anything; just report back what would have been done")
	cmd.Flags().BoolVar(&opts.interactive, "interactive", false, "Select interactively which containers to update")
	cmd.Flags().BoolVarP(&opts.force, "force", "f", false, "Disregard locks and container image filters (has no effect when used with --all or --update-all-images)")
	cmd.Flags().BoolVarP(&opts.watch, "watch", "w", false, "Watch rollout progress during release")

	// Deprecated
	cmd.Flags().StringSliceVarP(&opts.services, "service", "s", []string{}, "Service to release")
	cmd.Flags().MarkHidden("service")

	return cmd
}

func (opts *controllerReleaseOpts) RunE(cmd *cobra.Command, args []string) error {
	if len(opts.services) > 0 {
		return errorServiceFlagDeprecated
	}

	if len(args) != 0 {
		return errorWantedNoArgs
	}

	if err := checkExactlyOne("--update-image=<image> or --update-all-images", opts.image != "", opts.allImages); err != nil {
		return err
	}

	switch {
	case len(opts.controllers) <= 0 && !opts.allControllers:
		return newUsageError("please supply either --all, or at least one --controller=<controller>")
	case opts.force && opts.allControllers && opts.allImages:
		return newUsageError("--force has no effect when used with --all and --update-all-images")
	case opts.force && opts.allControllers:
		fmt.Fprintf(cmd.OutOrStderr(), "Warning: --force will not ignore locked controllers when used with --all\n")
	case opts.force && opts.allImages:
		fmt.Fprintf(cmd.OutOrStderr(), "Warning: --force will not ignore container image tags when used with --update-all-images\n")
	}

	var controllers []update.ResourceSpec

	if opts.allControllers {
		controllers = []update.ResourceSpec{update.ResourceSpecAll}
	} else {
		for _, controller := range opts.controllers {
			id, err := flux.ParseResourceIDOptionalNamespace(opts.namespace, controller)
			if err != nil {
				return err
			}
			controllers = append(controllers, update.MakeResourceSpec(id))
		}
	}

	var (
		image update.ImageSpec
		err   error
	)
	switch {
	case opts.image != "":
		image, err = update.ParseImageSpec(opts.image)
		if err != nil {
			return err
		}
	case opts.allImages:
		image = update.ImageSpecLatest
	}

	var kind update.ReleaseKind = update.ReleaseKindExecute
	if opts.dryRun || opts.interactive {
		kind = update.ReleaseKindPlan
	}

	var excludes []flux.ResourceID
	for _, exclude := range opts.exclude {
		s, err := flux.ParseResourceIDOptionalNamespace(opts.namespace, exclude)
		if err != nil {
			return err
		}
		excludes = append(excludes, s)
	}

	if kind == update.ReleaseKindPlan {
		fmt.Fprintf(cmd.OutOrStderr(), "Submitting dry-run release ...\n")
	} else {
		fmt.Fprintf(cmd.OutOrStderr(), "Submitting release ...\n")
	}

	ctx := context.Background()
	spec := update.ReleaseImageSpec{
		ServiceSpecs: controllers,
		ImageSpec:    image,
		Kind:         kind,
		Excludes:     excludes,
		Force:        opts.force,
	}
	jobID, err := opts.API.UpdateManifests(ctx, update.Spec{
		Type:  update.Images,
		Cause: opts.cause,
		Spec:  spec,
	})
	if err != nil {
		return err
	}

	result, err := awaitJob(ctx, opts.API, jobID)
	if err != nil {
		return err
	}
	if opts.interactive {
		spec, err := promptSpec(cmd.OutOrStdout(), result, opts.verbosity)
		spec.Force = opts.force
		if err != nil {
			fmt.Fprintln(cmd.OutOrStderr(), err.Error())
			return nil
		}

		fmt.Fprintf(cmd.OutOrStderr(), "Submitting selected release ...\n")
		jobID, err = opts.API.UpdateManifests(ctx, update.Spec{
			Type:  update.Containers,
			Cause: opts.cause,
			Spec:  spec,
		})
		if err != nil {
			fmt.Fprintln(cmd.OutOrStderr(), err.Error())
			return nil
		}

		opts.dryRun = false
	}

	if !opts.watch {
		return await(ctx, cmd.OutOrStdout(), cmd.OutOrStderr(), opts.API, jobID, !opts.dryRun, opts.verbosity)
	}

	fmt.Fprintf(cmd.OutOrStderr(), "Monitoring rollout ...\n")
	for {
		completed := 0
		services, err := opts.API.ListServicesWithOptions(ctx, v11.ListServicesOptions{Services: result.Result.AffectedResources()})
		if err != nil {
			return err
		}

		for _, service := range services {
			writeRolloutStatus(service, opts.verbosity)

			if service.Status == cluster.StatusReady {
				completed++
			}

			if service.Rollout.Messages != nil {
				fmt.Fprintf(cmd.OutOrStderr(), "There was a problem releasing %s:\n", service.ID)
				for _, msg := range service.Rollout.Messages {
					fmt.Fprintf(cmd.OutOrStderr(), "%s\n", msg)
				}
				return nil
			}
		}

		if completed == len(services) {
			fmt.Fprintf(cmd.OutOrStderr(), "All controllers ready.\n")
			return nil
		}

		time.Sleep(2000 * time.Millisecond)
	}
}

func writeRolloutStatus(service v6.ControllerStatus, verbosity int) {
	w := newTabwriter()
	fmt.Fprintf(w, "CONTROLLER\tCONTAINER\tIMAGE\tRELEASE\tREPLICAS\n")

	if len(service.Containers) > 0 {
		c := service.Containers[0]
		fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%d/%d", service.ID, c.Name, c.Current.ID, service.Status, service.Rollout.Ready, service.Rollout.Desired)
		if verbosity > 0 {
			fmt.Fprintf(w, " (%d outdated, %d updated)", service.Rollout.Outdated, service.Rollout.Updated)
		}
		fmt.Fprintf(w, "\n")
		for _, c := range service.Containers[1:] {
			fmt.Fprintf(w, "\t%s\t%s\t\t\n", c.Name, c.Current.ID)
		}
	} else {
		fmt.Fprintf(w, "%s\t\t\t%s\t%d/%d", service.ID, service.Status, service.Rollout.Ready, service.Rollout.Desired)
		if verbosity > 0 {
			fmt.Fprintf(w, " (%d outdated, %d updated)", service.Rollout.Outdated, service.Rollout.Updated)
		}
		fmt.Fprintf(w, "\n")
	}
	fmt.Fprintln(w)
	w.Flush()
}

func promptSpec(out io.Writer, result job.Result, verbosity int) (update.ReleaseContainersSpec, error) {
	menu := update.NewMenu(out, result.Result, verbosity)
	containerSpecs, err := menu.Run()
	return update.ReleaseContainersSpec{
		Kind:           update.ReleaseKindExecute,
		ContainerSpecs: containerSpecs,
		SkipMismatches: false,
	}, err
}
<|endoftext|>


<commit_before>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/scanner"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"

	"golang.org/x/tools/imports"
)

var (
	// main operation modes
	list    = flag.Bool("l", false, "list files whose formatting differs from goimport's")
	write   = flag.Bool("w", false, "write result to (source) file instead of stdout")
	doDiff  = flag.Bool("d", false, "display diffs instead of rewriting files")
	srcdir  = flag.String("srcdir", "", "choose imports as if source code is from `dir`. When operating on a single file, dir may instead be the complete file name.")
	verbose bool // verbose logging

	cpuProfile     = flag.String("cpuprofile", "", "CPU profile output")
	memProfile     = flag.String("memprofile", "", "memory profile output")
	memProfileRate = flag.Int("memrate", 0, "if > 0, sets runtime.MemProfileRate")

	options = &imports.Options{
		TabWidth:  8,
		TabIndent: true,
		Comments:  true,
		Fragment:  true,
	}
	exitCode = 0
)

func init() {
	flag.BoolVar(&options.AllErrors, "e", false, "report all errors (not just the first 10 on different lines)")
	flag.StringVar(&imports.LocalPrefix, "local", "", "put imports beginning with this string after 3rd-party packages")
}

func report(err error) {
	scanner.PrintError(os.Stderr, err)
	exitCode = 2
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: goimports [flags] [path ...]\n")
	flag.PrintDefaults()
	os.Exit(2)
}

func isGoFile(f os.FileInfo) bool {
	// ignore non-Go files
	name := f.Name()
	return !f.IsDir() && !strings.HasPrefix(name, ".") && strings.HasSuffix(name, ".go")
}

// argumentType is which mode goimports was invoked as.
type argumentType int

const (
	// fromStdin means the user is piping their source into goimports.
	fromStdin argumentType = iota

	// singleArg is the common case from editors, when goimports is run on
	// a single file.
	singleArg

	// multipleArg is when the user ran "goimports file1.go file2.go"
	// or ran goimports on a directory tree.
	multipleArg
)

func processFile(filename string, in io.Reader, out io.Writer, argType argumentType) error {
	opt := options
	if argType == fromStdin {
		nopt := *options
		nopt.Fragment = true
		opt = &nopt
	}

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

	src, err := ioutil.ReadAll(in)
	if err != nil {
		return err
	}

	target := filename
	if *srcdir != "" {
		// Determine whether the provided -srcdirc is a directory or file
		// and then use it to override the target.
		//
		// See https://github.com/dominikh/go-mode.el/issues/146
		if isFile(*srcdir) {
			if argType == multipleArg {
				return errors.New("-srcdir value can't be a file when passing multiple arguments or when walking directories")
			}
			target = *srcdir
		} else if argType == singleArg && strings.HasSuffix(*srcdir, ".go") && !isDir(*srcdir) {
			// For a file which doesn't exist on disk yet, but might shortly.
			// e.g. user in editor opens $DIR/newfile.go and newfile.go doesn't yet exist on disk.
			// The goimports on-save hook writes the buffer to a temp file
			// first and runs goimports before the actual save to newfile.go.
			// The editor's buffer is named "newfile.go" so that is passed to goimports as:
			//      goimports -srcdir=/gopath/src/pkg/newfile.go /tmp/gofmtXXXXXXXX.go
			// and then the editor reloads the result from the tmp file and writes
			// it to newfile.go.
			target = *srcdir
		} else {
			// Pretend that file is from *srcdir in order to decide
			// visible imports correctly.
			target = filepath.Join(*srcdir, filepath.Base(filename))
		}
	}

	res, err := imports.Process(target, src, opt)
	if err != nil {
		return err
	}

	if !bytes.Equal(src, res) {
		// formatting has changed
		if *list {
			fmt.Fprintln(out, filename)
		}
		if *write {
			err = ioutil.WriteFile(filename, res, 0)
			if err != nil {
				return err
			}
		}
		if *doDiff {
			data, err := diff(src, res)
			if err != nil {
				return fmt.Errorf("computing diff: %s", err)
			}
			fmt.Printf("diff %s gofmt/%s\n", filename, filename)
			out.Write(data)
		}
	}

	if !*list && !*write && !*doDiff {
		_, err = out.Write(res)
	}

	return err
}

func visitFile(path string, f os.FileInfo, err error) error {
	if err == nil && isGoFile(f) {
		err = processFile(path, nil, os.Stdout, multipleArg)
	}
	if err != nil {
		report(err)
	}
	return nil
}

func walkDir(path string) {
	filepath.Walk(path, visitFile)
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	// call gofmtMain in a separate function
	// so that it can use defer and have them
	// run before the exit.
	gofmtMain()
	os.Exit(exitCode)
}

// parseFlags parses command line flags and returns the paths to process.
// It's a var so that custom implementations can replace it in other files.
var parseFlags = func() []string {
	flag.BoolVar(&verbose, "v", false, "verbose logging")

	flag.Parse()
	return flag.Args()
}

func bufferedFileWriter(dest string) (w io.Writer, close func()) {
	f, err := os.Create(dest)
	if err != nil {
		log.Fatal(err)
	}
	bw := bufio.NewWriter(f)
	return bw, func() {
		if err := bw.Flush(); err != nil {
			log.Fatalf("error flushing %v: %v", dest, err)
		}
		if err := f.Close(); err != nil {
			log.Fatal(err)
		}
	}
}

func gofmtMain() {
	flag.Usage = usage
	paths := parseFlags()

	if *cpuProfile != "" {
		bw, flush := bufferedFileWriter(*cpuProfile)
		pprof.StartCPUProfile(bw)
		defer flush()
		defer pprof.StopCPUProfile()
	}
	// doTrace is a conditionally compiled wrapper around runtime/trace. It is
	// used to allow goimports to compile under gccgo, which does not support
	// runtime/trace. See https://golang.org/issue/15544.
	defer doTrace()()
	if *memProfileRate > 0 {
		runtime.MemProfileRate = *memProfileRate
		bw, flush := bufferedFileWriter(*memProfile)
		defer func() {
			runtime.GC() // materialize all statistics
			if err := pprof.WriteHeapProfile(bw); err != nil {
				log.Fatal(err)
			}
			flush()
		}()
	}

	if verbose {
		log.SetFlags(log.LstdFlags | log.Lmicroseconds)
		imports.Debug = true
	}
	if options.TabWidth < 0 {
		fmt.Fprintf(os.Stderr, "negative tabwidth %d\n", options.TabWidth)
		exitCode = 2
		return
	}

	if len(paths) == 0 {
		if err := processFile("<standard input>", os.Stdin, os.Stdout, fromStdin); err != nil {
			report(err)
		}
		return
	}

	argType := singleArg
	if len(paths) > 1 {
		argType = multipleArg
	}

	for _, path := range paths {
		switch dir, err := os.Stat(path); {
		case err != nil:
			report(err)
		case dir.IsDir():
			walkDir(path)
		default:
			if err := processFile(path, nil, os.Stdout, argType); err != nil {
				report(err)
			}
		}
	}
}

func diff(b1, b2 []byte) (data []byte, err error) {
	f1, err := ioutil.TempFile("", "gofmt")
	if err != nil {
		return
	}
	defer os.Remove(f1.Name())
	defer f1.Close()

	f2, err := ioutil.TempFile("", "gofmt")
	if err != nil {
		return
	}
	defer os.Remove(f2.Name())
	defer f2.Close()

	f1.Write(b1)
	f2.Write(b2)

	data, err = exec.Command("diff", "-u", f1.Name(), f2.Name()).CombinedOutput()
	if len(data) > 0 {
		// diff exits with a non-zero status when the files don't match.
		// Ignore that failure as long as we get output.
		err = nil
	}
	return
}

// isFile reports whether name is a file.
func isFile(name string) bool {
	fi, err := os.Stat(name)
	return err == nil && fi.Mode().IsRegular()
}

// isDir reports whether name is a directory.
func isDir(name string) bool {
	fi, err := os.Stat(name)
	return err == nil && fi.IsDir()
}
<commit_msg>cmd/goimports: use actual filename in goimports -d output<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/scanner"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"

	"golang.org/x/tools/imports"
)

var (
	// main operation modes
	list    = flag.Bool("l", false, "list files whose formatting differs from goimport's")
	write   = flag.Bool("w", false, "write result to (source) file instead of stdout")
	doDiff  = flag.Bool("d", false, "display diffs instead of rewriting files")
	srcdir  = flag.String("srcdir", "", "choose imports as if source code is from `dir`. When operating on a single file, dir may instead be the complete file name.")
	verbose bool // verbose logging

	cpuProfile     = flag.String("cpuprofile", "", "CPU profile output")
	memProfile     = flag.String("memprofile", "", "memory profile output")
	memProfileRate = flag.Int("memrate", 0, "if > 0, sets runtime.MemProfileRate")

	options = &imports.Options{
		TabWidth:  8,
		TabIndent: true,
		Comments:  true,
		Fragment:  true,
	}
	exitCode = 0
)

func init() {
	flag.BoolVar(&options.AllErrors, "e", false, "report all errors (not just the first 10 on different lines)")
	flag.StringVar(&imports.LocalPrefix, "local", "", "put imports beginning with this string after 3rd-party packages")
}

func report(err error) {
	scanner.PrintError(os.Stderr, err)
	exitCode = 2
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: goimports [flags] [path ...]\n")
	flag.PrintDefaults()
	os.Exit(2)
}

func isGoFile(f os.FileInfo) bool {
	// ignore non-Go files
	name := f.Name()
	return !f.IsDir() && !strings.HasPrefix(name, ".") && strings.HasSuffix(name, ".go")
}

// argumentType is which mode goimports was invoked as.
type argumentType int

const (
	// fromStdin means the user is piping their source into goimports.
	fromStdin argumentType = iota

	// singleArg is the common case from editors, when goimports is run on
	// a single file.
	singleArg

	// multipleArg is when the user ran "goimports file1.go file2.go"
	// or ran goimports on a directory tree.
	multipleArg
)

func processFile(filename string, in io.Reader, out io.Writer, argType argumentType) error {
	opt := options
	if argType == fromStdin {
		nopt := *options
		nopt.Fragment = true
		opt = &nopt
	}

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

	src, err := ioutil.ReadAll(in)
	if err != nil {
		return err
	}

	target := filename
	if *srcdir != "" {
		// Determine whether the provided -srcdirc is a directory or file
		// and then use it to override the target.
		//
		// See https://github.com/dominikh/go-mode.el/issues/146
		if isFile(*srcdir) {
			if argType == multipleArg {
				return errors.New("-srcdir value can't be a file when passing multiple arguments or when walking directories")
			}
			target = *srcdir
		} else if argType == singleArg && strings.HasSuffix(*srcdir, ".go") && !isDir(*srcdir) {
			// For a file which doesn't exist on disk yet, but might shortly.
			// e.g. user in editor opens $DIR/newfile.go and newfile.go doesn't yet exist on disk.
			// The goimports on-save hook writes the buffer to a temp file
			// first and runs goimports before the actual save to newfile.go.
			// The editor's buffer is named "newfile.go" so that is passed to goimports as:
			//      goimports -srcdir=/gopath/src/pkg/newfile.go /tmp/gofmtXXXXXXXX.go
			// and then the editor reloads the result from the tmp file and writes
			// it to newfile.go.
			target = *srcdir
		} else {
			// Pretend that file is from *srcdir in order to decide
			// visible imports correctly.
			target = filepath.Join(*srcdir, filepath.Base(filename))
		}
	}

	res, err := imports.Process(target, src, opt)
	if err != nil {
		return err
	}

	if !bytes.Equal(src, res) {
		// formatting has changed
		if *list {
			fmt.Fprintln(out, filename)
		}
		if *write {
			err = ioutil.WriteFile(filename, res, 0)
			if err != nil {
				return err
			}
		}
		if *doDiff {
			data, err := diff(src, res, filename)
			if err != nil {
				return fmt.Errorf("computing diff: %s", err)
			}
			fmt.Printf("diff -u %s %s\n", filepath.ToSlash(filename+".orig"), filepath.ToSlash(filename))
			out.Write(data)
		}
	}

	if !*list && !*write && !*doDiff {
		_, err = out.Write(res)
	}

	return err
}

func visitFile(path string, f os.FileInfo, err error) error {
	if err == nil && isGoFile(f) {
		err = processFile(path, nil, os.Stdout, multipleArg)
	}
	if err != nil {
		report(err)
	}
	return nil
}

func walkDir(path string) {
	filepath.Walk(path, visitFile)
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	// call gofmtMain in a separate function
	// so that it can use defer and have them
	// run before the exit.
	gofmtMain()
	os.Exit(exitCode)
}

// parseFlags parses command line flags and returns the paths to process.
// It's a var so that custom implementations can replace it in other files.
var parseFlags = func() []string {
	flag.BoolVar(&verbose, "v", false, "verbose logging")

	flag.Parse()
	return flag.Args()
}

func bufferedFileWriter(dest string) (w io.Writer, close func()) {
	f, err := os.Create(dest)
	if err != nil {
		log.Fatal(err)
	}
	bw := bufio.NewWriter(f)
	return bw, func() {
		if err := bw.Flush(); err != nil {
			log.Fatalf("error flushing %v: %v", dest, err)
		}
		if err := f.Close(); err != nil {
			log.Fatal(err)
		}
	}
}

func gofmtMain() {
	flag.Usage = usage
	paths := parseFlags()

	if *cpuProfile != "" {
		bw, flush := bufferedFileWriter(*cpuProfile)
		pprof.StartCPUProfile(bw)
		defer flush()
		defer pprof.StopCPUProfile()
	}
	// doTrace is a conditionally compiled wrapper around runtime/trace. It is
	// used to allow goimports to compile under gccgo, which does not support
	// runtime/trace. See https://golang.org/issue/15544.
	defer doTrace()()
	if *memProfileRate > 0 {
		runtime.MemProfileRate = *memProfileRate
		bw, flush := bufferedFileWriter(*memProfile)
		defer func() {
			runtime.GC() // materialize all statistics
			if err := pprof.WriteHeapProfile(bw); err != nil {
				log.Fatal(err)
			}
			flush()
		}()
	}

	if verbose {
		log.SetFlags(log.LstdFlags | log.Lmicroseconds)
		imports.Debug = true
	}
	if options.TabWidth < 0 {
		fmt.Fprintf(os.Stderr, "negative tabwidth %d\n", options.TabWidth)
		exitCode = 2
		return
	}

	if len(paths) == 0 {
		if err := processFile("<standard input>", os.Stdin, os.Stdout, fromStdin); err != nil {
			report(err)
		}
		return
	}

	argType := singleArg
	if len(paths) > 1 {
		argType = multipleArg
	}

	for _, path := range paths {
		switch dir, err := os.Stat(path); {
		case err != nil:
			report(err)
		case dir.IsDir():
			walkDir(path)
		default:
			if err := processFile(path, nil, os.Stdout, argType); err != nil {
				report(err)
			}
		}
	}
}

func diff(b1, b2 []byte, filename string) (data []byte, err error) {
	f1, err := ioutil.TempFile("", "gofmt")
	if err != nil {
		return
	}
	defer os.Remove(f1.Name())
	defer f1.Close()

	f2, err := ioutil.TempFile("", "gofmt")
	if err != nil {
		return
	}
	defer os.Remove(f2.Name())
	defer f2.Close()

	f1.Write(b1)
	f2.Write(b2)

	cmd := "diff"
	if runtime.GOOS == "plan9" {
		cmd = "/bin/ape/diff"
	}

	data, err = exec.Command(cmd, "-u", f1.Name(), f2.Name()).CombinedOutput()
	if len(data) > 0 {
		// diff exits with a non-zero status when the files don't match.
		// Ignore that failure as long as we get output.
		err = nil

		data, err = replaceTempFilename(data, filename)
	}
	return
}

// replaceTempFilename replaces temporary filenames in diff with actual one.
//
// --- /tmp/gofmt316145376	2017-02-03 19:13:00.280468375 -0500
// +++ /tmp/gofmt617882815	2017-02-03 19:13:00.280468375 -0500
// ...
// ->
// --- path/to/file.go.orig	2017-02-03 19:13:00.280468375 -0500
// +++ path/to/file.go	2017-02-03 19:13:00.280468375 -0500
// ...
func replaceTempFilename(diff []byte, filename string) ([]byte, error) {
	bs := bytes.SplitN(diff, []byte{'\n'}, 3)
	if len(bs) < 3 {
		return nil, fmt.Errorf("got unexpected diff for %s", filename)
	}
	// Preserve timestamps.
	var t0, t1 []byte
	if i := bytes.LastIndexByte(bs[0], '\t'); i != -1 {
		t0 = bs[0][i:]
	}
	if i := bytes.LastIndexByte(bs[1], '\t'); i != -1 {
		t1 = bs[1][i:]
	}
	// Always print filepath with slash separator.
	f := filepath.ToSlash(filename)
	bs[0] = []byte(fmt.Sprintf("--- %s%s", f+".orig", t0))
	bs[1] = []byte(fmt.Sprintf("+++ %s%s", f, t1))
	return bytes.Join(bs, []byte{'\n'}), nil
}

// isFile reports whether name is a file.
func isFile(name string) bool {
	fi, err := os.Stat(name)
	return err == nil && fi.Mode().IsRegular()
}

// isDir reports whether name is a directory.
func isDir(name string) bool {
	fi, err := os.Stat(name)
	return err == nil && fi.IsDir()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"html/template"
	"math/big"
	"net/http"
	"strings"
	"time"

	"github.com/go-martini/martini"
	"github.com/howeyc/ledger"
	"github.com/lucasb-eyer/go-colorful"
)

func getRangeAndPeriod(dateRange, dateFreq string) (start, end time.Time, period ledger.Period) {
	switch dateFreq {
	case "Monthly":
		period = ledger.PeriodMonth
	case "Quarterly":
		period = ledger.PeriodQuarter
	case "Yearly":
		period = ledger.PeriodYear
	default:
		period = ledger.PeriodMonth
	}

	currentTime := time.Now()
	switch dateRange {
	case "All Time":
		end = currentTime.Add(time.Hour * 24)
	case "YTD":
		start = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
		end = currentTime.Add(time.Hour * 24)
	case "Previous Year":
		start = time.Date(currentTime.Year()-1, time.January, 1, 0, 0, 0, 0, time.UTC)
		end = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
	case "Previous Month":
		start = time.Date(currentTime.Year(), currentTime.Month()-1, 1, 0, 0, 0, 0, time.UTC)
		end = time.Date(currentTime.Year(), currentTime.Month(), 1, 0, 0, 0, 0, time.UTC)
	case "Current Month":
		start = time.Date(currentTime.Year(), currentTime.Month(), 1, 0, 0, 0, 0, time.UTC)
		end = currentTime.Add(time.Hour * 24)
	case "Current Quarter":
		switch currentTime.Month() {
		case time.January, time.February, time.March:
			start = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
		case time.April, time.May, time.June:
			start = time.Date(currentTime.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
		case time.July, time.August, time.September:
			start = time.Date(currentTime.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.October, 1, 0, 0, 0, 0, time.UTC)
		case time.October, time.November, time.December:
			start = time.Date(currentTime.Year(), time.October, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year()+1, time.January, 1, 0, 0, 0, 0, time.UTC)
		}
	case "Previous Quarter":
		switch currentTime.Month() {
		case time.January, time.February, time.March:
			start = time.Date(currentTime.Year()-1, time.October, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
		case time.April, time.May, time.June:
			start = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
		case time.July, time.August, time.September:
			start = time.Date(currentTime.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
		case time.October, time.November, time.December:
			start = time.Date(currentTime.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.October, 1, 0, 0, 0, 0, time.UTC)
		}
	}

	return
}

// getAccounts will return the accounts that match accountNeedle.
// If accountNeedle contains no wildcards (*), only case-sensitive matchs are returned.
// In the case of wildcards:
//	* matches any account that shares the prefix, and has the same depth as the *.
//	** matches any account that shares the prefix, at any depth.
func getAccounts(accountNeedle string, accountsHaystack []*ledger.Account) (results []*ledger.Account) {
	needleDepth := len(strings.Split(accountNeedle, ":"))

	if dblstarIdx := strings.Index(accountNeedle, "**"); dblstarIdx != -1 {
		prefixNeedle := accountNeedle[:dblstarIdx]
		for _, hay := range accountsHaystack {
			if strings.HasPrefix(hay.Name, prefixNeedle) {
				results = append(results, hay)
			}
		}
	} else if starIdx := strings.Index(accountNeedle, "*"); starIdx != -1 {
		prefixNeedle := accountNeedle[:starIdx]
		for _, hay := range accountsHaystack {
			hayDepth := len(strings.Split(hay.Name, ":"))
			if strings.HasPrefix(hay.Name, prefixNeedle) && hayDepth == needleDepth {
				results = append(results, hay)
			}
		}
	} else {
		for _, hay := range accountsHaystack {
			if hay.Name == accountNeedle {
				results = append(results, hay)
			}
		}
	}

	return
}

func calcBalances(calcAccts []calculatedAccount, balances []*ledger.Account) (results []*ledger.Account) {
	accVals := make(map[string]*big.Rat)
	for _, calcAccount := range calcAccts {
		for _, bal := range balances {
			for _, acctOp := range calcAccount.AccountOperations {
				if acctOp.Name == bal.Name {
					fval := big.NewRat(1, 1).Abs(bal.Balance)
					aval, found := accVals[calcAccount.Name]
					if !found {
						aval = big.NewRat(0, 1)
					}
					switch acctOp.Operation {
					case "+":
						aval.Add(aval, fval)
					case "-":
						aval.Sub(aval, fval)
					}
					accVals[calcAccount.Name] = aval
				}
			}
		}
	}

	for _, calcAccount := range calcAccts {
		results = append(results, &ledger.Account{Name: calcAccount.Name, Balance: accVals[calcAccount.Name]})
	}

	return
}

func reportHandler(w http.ResponseWriter, r *http.Request, params martini.Params) {
	reportName := params["reportName"]

	trans, terr := getTransactions()
	if terr != nil {
		http.Error(w, terr.Error(), 500)
		return
	}

	var rConf reportConfig
	for _, reportConf := range reportConfigData.Reports {
		if reportConf.Name == reportName {
			rConf = reportConf
		}
	}
	rStart, rEnd, rPeriod := getRangeAndPeriod(rConf.DateRange, rConf.DateFreq)

	trans = ledger.TransactionsInDateRange(trans, rStart, rEnd)

	var rtrans []*ledger.Transaction
	for _, tran := range trans {
		include := true
		for _, accChange := range tran.AccountChanges {
			for _, excludeName := range rConf.ExcludeAccountTrans {
				if strings.Contains(accChange.Name, excludeName) {
					include = false
				}
			}
		}

		if include {
			rtrans = append(rtrans, tran)
		}
	}

	balances := ledger.GetBalances(rtrans, []string{})
	var initialAccounts []*ledger.Account
	for _, confAccount := range rConf.Accounts {
		initialAccounts = append(initialAccounts, getAccounts(confAccount, balances)...)
	}
	initialAccounts = append(initialAccounts, calcBalances(rConf.CalculatedAccounts, balances)...)
	var reportSummaryAccounts []*ledger.Account
	for _, account := range initialAccounts {
		include := true
		for _, excludeName := range rConf.ExcludeAccountsSummary {
			if strings.Contains(account.Name, excludeName) {
				include = false
			}
		}

		if include {
			reportSummaryAccounts = append(reportSummaryAccounts, account)
		}
	}

	// Filter report to only show transactions that are for the accounts in the summary of the report
	var vtrans []*ledger.Transaction
	for _, trans := range rtrans {
		include := false
		for _, accChange := range trans.AccountChanges {
			for _, account := range reportSummaryAccounts {
				if strings.Contains(accChange.Name, account.Name) {
					include = true
				}
			}
		}
		if include {
			vtrans = append(vtrans, trans)
		}
	}

	colorPalette, cerr := colorful.HappyPalette(len(reportSummaryAccounts))
	if cerr != nil {
		http.Error(w, cerr.Error(), 500)
		return
	}
	colorBlack := colorful.Color{R: 1, G: 1, B: 1}

	switch rConf.Chart {
	case "pie":
		type pieAccount struct {
			Name      string
			Balance   *big.Rat
			Color     string
			Highlight string
		}

		var values []pieAccount

		type pieColor struct {
			Color     string
			Highlight string
		}

		for colorIdx, account := range reportSummaryAccounts {
			accName := account.Name
			value := big.NewRat(1, 1).Set(account.Balance)
			values = append(values, pieAccount{Name: accName, Balance: value,
				Color:     colorPalette[colorIdx].Hex(),
				Highlight: colorPalette[colorIdx].BlendRgb(colorBlack, 0.1).Hex()})
		}

		type piePageData struct {
			pageData
			ReportName           string
			RangeStart, RangeEnd time.Time
			ChartAccounts        []pieAccount
		}

		var pData piePageData
		pData.Reports = reportConfigData.Reports
		pData.Transactions = vtrans
		pData.ChartAccounts = values
		pData.RangeStart = rStart
		pData.RangeEnd = rEnd
		pData.ReportName = reportName

		funcMap := template.FuncMap{
			"abbrev": abbrev,
		}

		t, err := parseAssetsWithFunc(funcMap, "templates/template.piechart.html", "templates/template.nav.html")
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		err = t.Execute(w, pData)
		if err != nil {
			http.Error(w, err.Error(), 500)
		}
	case "line", "bar", "stackedbar":
		type lineData struct {
			AccountName string
			RGBColor    string
			Values      []*big.Rat
		}
		type linePageData struct {
			pageData
			ReportName           string
			RangeStart, RangeEnd time.Time
			ChartType            string
			Labels               []string
			DataSets             []lineData
		}
		var lData linePageData
		lData.Reports = reportConfigData.Reports
		lData.ReportName = reportName

		for colorIdx, repAccount := range reportSummaryAccounts {
			r, g, b := colorPalette[colorIdx].RGB255()
			lData.DataSets = append(lData.DataSets,
				lineData{AccountName: repAccount.Name,
					RGBColor: fmt.Sprintf("%d, %d, %d", r, g, b)})
		}

		var rType ledger.RangeType
		switch rConf.Chart {
		case "line":
			rType = ledger.RangeSnapshot
			lData.ChartType = "Line"
		case "bar":
			rType = ledger.RangePartition
			lData.ChartType = "Bar"
		case "stackedbar":
			rType = ledger.RangePartition
			lData.ChartType = "StackedBar"
		}

		lData.Transactions = vtrans

		rangeBalances := ledger.BalancesByPeriod(rtrans, rPeriod, rType)
		for _, rb := range rangeBalances {
			if lData.RangeStart.IsZero() {
				lData.RangeStart = rb.Start
			}
			lData.RangeEnd = rb.End
			lData.Labels = append(lData.Labels, rb.End.Format("2006-01-02"))

			accVals := make(map[string]*big.Rat)
			for _, confAccount := range rConf.Accounts {
				for _, freqAccount := range getAccounts(confAccount, rb.Balances) {
					accVals[freqAccount.Name] = big.NewRat(1, 1).Abs(freqAccount.Balance)
				}
			}

			for _, calcAccount := range calcBalances(rConf.CalculatedAccounts, rb.Balances) {
				accVals[calcAccount.Name] = calcAccount.Balance
			}

			for dIdx := range lData.DataSets {
				aval, afound := accVals[lData.DataSets[dIdx].AccountName]
				if !afound {
					aval = big.NewRat(0, 1)
				}
				lData.DataSets[dIdx].Values = append(lData.DataSets[dIdx].Values, aval)
			}
		}

		t, err := parseAssets("templates/template.barlinechart.html", "templates/template.nav.html")
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		err = t.Execute(w, lData)
		if err != nil {
			http.Error(w, err.Error(), 500)
		}

	default:
		fmt.Fprint(w, "Unsupported chart type.")
	}
}
<commit_msg>small fix<commit_after>package main

import (
	"fmt"
	"html/template"
	"math/big"
	"net/http"
	"strings"
	"time"

	"github.com/go-martini/martini"
	"github.com/howeyc/ledger"
	"github.com/lucasb-eyer/go-colorful"
)

func getRangeAndPeriod(dateRange, dateFreq string) (start, end time.Time, period ledger.Period) {
	switch dateFreq {
	case "Monthly":
		period = ledger.PeriodMonth
	case "Quarterly":
		period = ledger.PeriodQuarter
	case "Yearly":
		period = ledger.PeriodYear
	default:
		period = ledger.PeriodMonth
	}

	currentTime := time.Now()
	switch dateRange {
	case "All Time":
		end = currentTime.Add(time.Hour * 24)
	case "YTD":
		start = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
		end = currentTime.Add(time.Hour * 24)
	case "Previous Year":
		start = time.Date(currentTime.Year()-1, time.January, 1, 0, 0, 0, 0, time.UTC)
		end = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
	case "Previous Month":
		start = time.Date(currentTime.Year(), currentTime.Month()-1, 1, 0, 0, 0, 0, time.UTC)
		end = time.Date(currentTime.Year(), currentTime.Month(), 1, 0, 0, 0, 0, time.UTC)
	case "Current Month":
		start = time.Date(currentTime.Year(), currentTime.Month(), 1, 0, 0, 0, 0, time.UTC)
		end = currentTime.Add(time.Hour * 24)
	case "Current Quarter":
		switch currentTime.Month() {
		case time.January, time.February, time.March:
			start = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
		case time.April, time.May, time.June:
			start = time.Date(currentTime.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
		case time.July, time.August, time.September:
			start = time.Date(currentTime.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.October, 1, 0, 0, 0, 0, time.UTC)
		case time.October, time.November, time.December:
			start = time.Date(currentTime.Year(), time.October, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year()+1, time.January, 1, 0, 0, 0, 0, time.UTC)
		}
	case "Previous Quarter":
		switch currentTime.Month() {
		case time.January, time.February, time.March:
			start = time.Date(currentTime.Year()-1, time.October, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
		case time.April, time.May, time.June:
			start = time.Date(currentTime.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
		case time.July, time.August, time.September:
			start = time.Date(currentTime.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
		case time.October, time.November, time.December:
			start = time.Date(currentTime.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
			end = time.Date(currentTime.Year(), time.October, 1, 0, 0, 0, 0, time.UTC)
		}
	}

	return
}

// getAccounts will return the accounts that match accountNeedle.
// If accountNeedle contains no wildcards (*), only case-sensitive matchs are returned.
// In the case of wildcards:
//	* matches any account that shares the prefix, and has the same depth as the *.
//	** matches any account that shares the prefix, at any depth.
func getAccounts(accountNeedle string, accountsHaystack []*ledger.Account) (results []*ledger.Account) {
	needleDepth := len(strings.Split(accountNeedle, ":"))

	if dblstarIdx := strings.Index(accountNeedle, "**"); dblstarIdx != -1 {
		prefixNeedle := accountNeedle[:dblstarIdx]
		for _, hay := range accountsHaystack {
			if strings.HasPrefix(hay.Name, prefixNeedle) {
				results = append(results, hay)
			}
		}
	} else if starIdx := strings.Index(accountNeedle, "*"); starIdx != -1 {
		prefixNeedle := accountNeedle[:starIdx]
		for _, hay := range accountsHaystack {
			hayDepth := len(strings.Split(hay.Name, ":"))
			if strings.HasPrefix(hay.Name, prefixNeedle) && hayDepth == needleDepth {
				results = append(results, hay)
			}
		}
	} else {
		for _, hay := range accountsHaystack {
			if hay.Name == accountNeedle {
				results = append(results, hay)
			}
		}
	}

	return
}

func calcBalances(calcAccts []calculatedAccount, balances []*ledger.Account) (results []*ledger.Account) {
	accVals := make(map[string]*big.Rat)
	for _, calcAccount := range calcAccts {
		for _, bal := range balances {
			for _, acctOp := range calcAccount.AccountOperations {
				if acctOp.Name == bal.Name {
					fval := big.NewRat(1, 1).Abs(bal.Balance)
					aval, found := accVals[calcAccount.Name]
					if !found {
						aval = big.NewRat(0, 1)
					}
					switch acctOp.Operation {
					case "+":
						aval.Add(aval, fval)
					case "-":
						aval.Sub(aval, fval)
					}
					accVals[calcAccount.Name] = aval
				}
			}
		}
	}

	for _, calcAccount := range calcAccts {
		results = append(results, &ledger.Account{Name: calcAccount.Name, Balance: accVals[calcAccount.Name]})
	}

	return
}

func reportHandler(w http.ResponseWriter, r *http.Request, params martini.Params) {
	reportName := params["reportName"]

	trans, terr := getTransactions()
	if terr != nil {
		http.Error(w, terr.Error(), 500)
		return
	}

	var rConf reportConfig
	for _, reportConf := range reportConfigData.Reports {
		if reportConf.Name == reportName {
			rConf = reportConf
		}
	}
	rStart, rEnd, rPeriod := getRangeAndPeriod(rConf.DateRange, rConf.DateFreq)

	trans = ledger.TransactionsInDateRange(trans, rStart, rEnd)

	var rtrans []*ledger.Transaction
	for _, tran := range trans {
		include := true
		for _, accChange := range tran.AccountChanges {
			for _, excludeName := range rConf.ExcludeAccountTrans {
				if strings.Contains(accChange.Name, excludeName) {
					include = false
				}
			}
		}

		if include {
			rtrans = append(rtrans, tran)
		}
	}

	balances := ledger.GetBalances(rtrans, []string{})
	var initialAccounts []*ledger.Account
	for _, confAccount := range rConf.Accounts {
		initialAccounts = append(initialAccounts, getAccounts(confAccount, balances)...)
	}
	initialAccounts = append(initialAccounts, calcBalances(rConf.CalculatedAccounts, balances)...)
	var reportSummaryAccounts []*ledger.Account
	for _, account := range initialAccounts {
		include := true
		for _, excludeName := range rConf.ExcludeAccountsSummary {
			if strings.Contains(account.Name, excludeName) {
				include = false
			}
		}

		if include {
			reportSummaryAccounts = append(reportSummaryAccounts, account)
		}
	}

	// Filter report to only show transactions that are for the accounts in the summary of the report
	var vtrans []*ledger.Transaction
	for _, trans := range rtrans {
		include := false
		for _, accChange := range trans.AccountChanges {
			for _, account := range reportSummaryAccounts {
				if strings.Contains(accChange.Name, account.Name) {
					include = true
				}
			}
		}
		if include {
			vtrans = append(vtrans, trans)
		}
	}

	colorPalette, cerr := colorful.HappyPalette(len(reportSummaryAccounts))
	if cerr != nil {
		http.Error(w, cerr.Error(), 500)
		return
	}
	colorBlack := colorful.Color{R: 1, G: 1, B: 1}

	switch rConf.Chart {
	case "pie":
		type pieAccount struct {
			Name      string
			Balance   *big.Rat
			Color     string
			Highlight string
		}

		var values []pieAccount

		type pieColor struct {
			Color     string
			Highlight string
		}

		for colorIdx, account := range reportSummaryAccounts {
			accName := account.Name
			value := big.NewRat(1, 1).Set(account.Balance)
			values = append(values, pieAccount{Name: accName, Balance: value,
				Color:     colorPalette[colorIdx].Hex(),
				Highlight: colorPalette[colorIdx].BlendRgb(colorBlack, 0.1).Hex()})
		}

		type piePageData struct {
			pageData
			ReportName           string
			RangeStart, RangeEnd time.Time
			ChartAccounts        []pieAccount
		}

		var pData piePageData
		pData.Reports = reportConfigData.Reports
		pData.Transactions = vtrans
		pData.ChartAccounts = values
		pData.RangeStart = rStart
		pData.RangeEnd = rEnd
		pData.ReportName = reportName

		funcMap := template.FuncMap{
			"abbrev": abbrev,
		}

		t, err := parseAssetsWithFunc(funcMap, "templates/template.piechart.html", "templates/template.nav.html")
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		err = t.Execute(w, pData)
		if err != nil {
			http.Error(w, err.Error(), 500)
		}
	case "line", "bar", "stackedbar":
		type lineData struct {
			AccountName string
			RGBColor    string
			Values      []*big.Rat
		}
		type linePageData struct {
			pageData
			ReportName           string
			RangeStart, RangeEnd time.Time
			ChartType            string
			Labels               []string
			DataSets             []lineData
		}
		var lData linePageData
		lData.Reports = reportConfigData.Reports
		lData.ReportName = reportName

		for colorIdx, repAccount := range reportSummaryAccounts {
			r, g, b := colorPalette[colorIdx].RGB255()
			lData.DataSets = append(lData.DataSets,
				lineData{AccountName: repAccount.Name,
					RGBColor: fmt.Sprintf("%d, %d, %d", r, g, b)})
		}

		var rType ledger.RangeType
		switch rConf.Chart {
		case "line":
			rType = ledger.RangeSnapshot
			lData.ChartType = "Line"
		case "bar":
			rType = ledger.RangePartition
			lData.ChartType = "Bar"
		case "stackedbar":
			rType = ledger.RangePartition
			lData.ChartType = "StackedBar"
		}

		lData.Transactions = vtrans

		rangeBalances := ledger.BalancesByPeriod(rtrans, rPeriod, rType)
		for _, rb := range rangeBalances {
			if lData.RangeStart.IsZero() {
				lData.RangeStart = rb.Start
			}
			lData.RangeEnd = rb.End
			lData.Labels = append(lData.Labels, rb.End.Format("2006-01-02"))

			accVals := make(map[string]*big.Rat)
			for _, confAccount := range rConf.Accounts {
				for _, freqAccount := range getAccounts(confAccount, rb.Balances) {
					accVals[freqAccount.Name] = big.NewRat(1, 1).Abs(freqAccount.Balance)
				}
			}

			for _, calcAccount := range calcBalances(rConf.CalculatedAccounts, rb.Balances) {
				accVals[calcAccount.Name] = calcAccount.Balance
			}

			for dIdx := range lData.DataSets {
				aval, afound := accVals[lData.DataSets[dIdx].AccountName]
				if !afound || aval == nil {
					aval = big.NewRat(0, 1)
				}
				lData.DataSets[dIdx].Values = append(lData.DataSets[dIdx].Values, aval)
			}
		}

		t, err := parseAssets("templates/template.barlinechart.html", "templates/template.nav.html")
		if err != nil {
			http.Error(w, err.Error(), 500)
			return
		}
		err = t.Execute(w, lData)
		if err != nil {
			http.Error(w, err.Error(), 500)
		}

	default:
		fmt.Fprint(w, "Unsupported chart type.")
	}
}
<|endoftext|>


<commit_before>// Copyright © 2017 Salim Alami Idrissi
//
// 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 cmd

import (
	"context"
	"fmt"
	"io"
	"log"
	"time"

	"github.com/celrenheit/sandglass/cmd/cmdcommon"

	"golang.org/x/sync/errgroup"

	"github.com/spf13/viper"

	"github.com/celrenheit/sandflake"
	"github.com/celrenheit/sandglass-grpc/go/sgproto"
	"github.com/spf13/cobra"
)

// consumeCmd represents the list command
var consumeCmd = &cobra.Command{
	Use:   "consume",
	Short: "",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) != 1 {
			log.Fatal("only one topic is allowed")
		}
		topic := args[0]

		consumerGroup := viper.GetString("consumer-group")
		consumerName := viper.GetString("consumer-name")
		if consumerName == "" {
			var gen sandflake.Generator
			consumerName = gen.Next().String()
		}

		var group errgroup.Group
		msgCh := make(chan *sgproto.Message)
		if cmd.Flag("partition").Changed {
			group.Go(func() error {
				consume(msgCh, topic, viper.GetString("partition"), consumerGroup, consumerName, viper.GetBool("follow"), viper.GetBool("ack"))
				return nil
			})
		} else {
			t, err := client.GetTopic(context.Background(), &sgproto.GetTopicParams{
				Name: topic,
			})
			if err != nil {
				panic(err)
			}

			for _, part := range t.Partitions {
				part := part
				group.Go(func() error {
					consume(msgCh, topic, part, consumerGroup, consumerName, viper.GetBool("follow"), viper.GetBool("ack"))
					return nil
				})
			}
		}

		go func() {
			group.Wait()
			close(msgCh)
		}()

		for msg := range msgCh {
			fmt.Println(string(msg.Value))
		}

	},
}

func init() {
	RootCmd.AddCommand(consumeCmd)

	consumeCmd.Flags().String("partition", "random", "Partition to produce in")
	consumeCmd.Flags().String("consumer-group", "sandctl", "Consumer group")
	consumeCmd.Flags().String("consumer-name", "", "Consumer name (default: random)")
	consumeCmd.Flags().Duration("poll-interval", 50*time.Millisecond, "Poll interval")
	consumeCmd.Flags().BoolP("follow", "f", false, "Consumer name (default: random)")
	consumeCmd.Flags().Bool("ack", true, "Ack messages (batching 10k messages)")

	cmdcommon.BindViper(consumeCmd.Flags(),
		"partition",
		"consumer-group",
		"consumer-name",
		"poll-interval",
		"follow",
		"ack",
	)
}

func consume(msgCh chan *sgproto.Message, topic, partition, group, name string, follow, ack bool) {
	ctx := context.Background()

FOLLOW:

	stream, err := client.ConsumeFromGroup(ctx, &sgproto.ConsumeFromGroupRequest{
		Topic:             topic,
		Partition:         partition,
		ConsumerGroupName: group,
		ConsumerName:      name,
	})
	if err != nil {
		panic(err)
	}

	ackFn := func(offsets []sandflake.ID) error {
		_, err := client.AcknowledgeMessages(context.Background(), &sgproto.MultiOffsetChangeRequest{
			Topic:         topic,
			Partition:     partition,
			ConsumerGroup: group,
			ConsumerName:  name,
			Offsets:       offsets,
		})
		return err
	}

	offsets := []sandflake.ID{}
	for {
		msg, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			panic(err)
		}

		msgCh <- msg

		if ack {
			offsets = append(offsets, msg.Offset)
		}

		if ack && len(offsets) == 1000 {
			err := ackFn(offsets)
			if err != nil {
				panic(err)
			}

			offsets = offsets[:0]
		}
	}

	if ack && len(offsets) > 0 {
		err := ackFn(offsets)
		if err != nil {
			panic(err)
		}
	}

	if follow {
		time.Sleep(viper.GetDuration("poll-interval"))
		goto FOLLOW
	}
}
<commit_msg>fix sandctl<commit_after>// Copyright © 2017 Salim Alami Idrissi
//
// 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 cmd

import (
	"context"
	"fmt"
	"io"
	"log"
	"time"

	"github.com/celrenheit/sandglass/cmd/cmdcommon"

	"golang.org/x/sync/errgroup"

	"github.com/spf13/viper"

	"github.com/celrenheit/sandflake"
	"github.com/celrenheit/sandglass-grpc/go/sgproto"
	"github.com/spf13/cobra"
)

// consumeCmd represents the list command
var consumeCmd = &cobra.Command{
	Use:   "consume",
	Short: "",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) != 1 {
			log.Fatal("only one topic is allowed")
		}
		topic := args[0]

		consumerGroup := viper.GetString("consumer-group")
		consumerName := viper.GetString("consumer-name")
		if consumerName == "" {
			var gen sandflake.Generator
			consumerName = gen.Next().String()
		}

		var group errgroup.Group
		msgCh := make(chan *sgproto.Message)
		if cmd.Flag("partition").Changed {
			group.Go(func() error {
				consume(msgCh, topic, viper.GetString("partition"), consumerGroup, consumerName, viper.GetBool("follow"), viper.GetBool("ack"))
				return nil
			})
		} else {
			t, err := client.GetTopic(context.Background(), &sgproto.GetTopicParams{
				Name: topic,
			})
			if err != nil {
				panic(err)
			}

			for _, part := range t.Partitions {
				part := part
				group.Go(func() error {
					consume(msgCh, topic, part, consumerGroup, consumerName, viper.GetBool("follow"), viper.GetBool("ack"))
					return nil
				})
			}
		}

		go func() {
			group.Wait()
			close(msgCh)
		}()

		for msg := range msgCh {
			fmt.Println(string(msg.Value))
		}

	},
}

func init() {
	RootCmd.AddCommand(consumeCmd)

	consumeCmd.Flags().String("partition", "random", "Partition to produce in")
	consumeCmd.Flags().String("consumer-group", "sandctl", "Consumer group")
	consumeCmd.Flags().String("consumer-name", "", "Consumer name (default: random)")
	consumeCmd.Flags().Duration("poll-interval", 50*time.Millisecond, "Poll interval")
	consumeCmd.Flags().BoolP("follow", "f", false, "Consumer name (default: random)")
	consumeCmd.Flags().Bool("ack", true, "Ack messages (batching 10k messages)")

	cmdcommon.BindViper(consumeCmd.Flags(),
		"partition",
		"consumer-group",
		"consumer-name",
		"poll-interval",
		"follow",
		"ack",
	)
}

func consume(msgCh chan *sgproto.Message, topic, partition, group, name string, follow, ack bool) {
	ctx := context.Background()

FOLLOW:

	stream, err := client.ConsumeFromGroup(ctx, &sgproto.ConsumeFromGroupRequest{
		Topic:             topic,
		Partition:         partition,
		ConsumerGroupName: group,
		ConsumerName:      name,
	})
	if err != nil {
		panic(err)
	}

	ackFn := func(offsets []sandflake.ID) error {
		_, err := client.Acknowledge(context.Background(), &sgproto.MarkRequest{
			Topic:         topic,
			Partition:     partition,
			ConsumerGroup: group,
			ConsumerName:  name,
			Offsets:       offsets,
		})
		return err
	}

	offsets := []sandflake.ID{}
	for {
		msg, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			panic(err)
		}

		msgCh <- msg

		if ack {
			offsets = append(offsets, msg.Offset)
		}

		if ack && len(offsets) == 1000 {
			err := ackFn(offsets)
			if err != nil {
				panic(err)
			}

			offsets = offsets[:0]
		}
	}

	if ack && len(offsets) > 0 {
		err := ackFn(offsets)
		if err != nil {
			panic(err)
		}
	}

	if follow {
		time.Sleep(viper.GetDuration("poll-interval"))
		goto FOLLOW
	}
}
<|endoftext|>