<commit_before>// Copyright 2016 The Upspin 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 (
	"flag"
	"io"
	"log"
	"os"
	"path/filepath"

	"upspin.io/errors"
	"upspin.io/path"
	"upspin.io/upspin"
)

func (s *State) cp(args ...string) {
	const help = `
Cp copies files into, out of, and within Upspin. If the final
argument is a directory, the files are placed inside it.  The other
arguments must not be directories unless the -R flag is set.

If the final argument is not a directory, cp requires exactly two
path names and copies the contents of the first to the second.
The -R flag requires that the final argument be a directory.

When copying from one Upspin path to another Upspin path, cp can be
very efficient, copying only the references to the data rather than
the data itself.
`
	fs := flag.NewFlagSet("cp", flag.ExitOnError)
	fs.Bool("v", false, "log each file as it is copied")
	fs.Bool("R", false, "recursively copy directories")
	s.parseFlags(fs, args, help, "cp [opts] file... file or cp [opts] file... directory")

	cs := &copyState{
		state:   s,
		flagSet: fs,
		recur:   boolFlag(fs, "R"),
		verbose: boolFlag(fs, "v"),
	}

	// Do all the glob processing here.
	// Special one-at-time glob processing because each item may be local or Upspin.
	var files []cpFile
	for _, file := range fs.Args() {
		files = append(files, cs.glob(file)...)
	}

	if len(files) < 2 {
		fs.Usage()
	}

	nSrc := len(files) - 1
	src, dest := files[:nSrc], files[nSrc]
	s.copyCommand(cs, src, dest)
}

type copyState struct {
	state   *State
	flagSet *flag.FlagSet // Used only to call Usage.
	verbose bool
	recur   bool
}

func (c *copyState) logf(format string, args ...interface{}) {
	if c.verbose {
		log.Printf(format, args...)
	}
}

// A cpFile is a glob-expanded file name and an indication of whether
// it resides on Upspin.
type cpFile struct {
	path     string
	isUpspin bool
}

var (
	errExist    = errors.E(errors.Exist)
	errNotExist = errors.E(errors.NotExist)
	errIsDir    = errors.E(errors.IsDir)
)

func (s *State) copyCommand(cs *copyState, srcFiles []cpFile, dstFile cpFile) {
	// TODO: Check for nugatory copies.
	if s.isDir(dstFile) {
		s.copyToDir(cs, srcFiles, dstFile)
		return
	}
	if len(srcFiles) != 1 {
		s.failf("copying multiple files but %s is not a directory", dstFile.path)
		cs.flagSet.Usage()
	}
	if cs.recur {
		s.failf("recursive copy requires that final argument (%s) be an existing directory", dstFile.path)
		cs.flagSet.Usage()
	}
	reader, err := s.open(srcFiles[0])
	if err != nil {
		s.fail(err)
	}
	s.copyToFile(cs, reader, srcFiles[0], dstFile)
}

// isDir reports whether the file is a directory either in Upspin
// or in the local file system.
func (s *State) isDir(cf cpFile) bool {
	if cf.isUpspin {
		entry, err := s.client.Lookup(upspin.PathName(cf.path), true)
		// Report the error here if it's anything odd, because otherwise
		// we'll report "not a directory" misleadingly.
		if err != nil && !errors.Match(errNotExist, err) {
			log.Printf("%q: %v", cf.path, err)
		}
		return err == nil && entry.IsDir()
	}
	// Not an Upspin name. Is it a local directory?
	info, err := os.Stat(cf.path)
	return err == nil && info.IsDir()
}

// open opens the file regardless of its location.
func (s *State) open(file cpFile) (io.ReadCloser, error) {
	if s.isDir(file) {
		return nil, errors.E(upspin.PathName(file.path), errors.IsDir)
	}
	if file.isUpspin {
		return s.client.Open(upspin.PathName(file.path))
	}
	return os.Open(file.path)
}

// create creates the file regardless of its location.
func (s *State) create(file cpFile) (io.WriteCloser, error) {
	if file.isUpspin {
		fd, err := s.client.Create(upspin.PathName(file.path))
		return fd, err
	}
	fd, err := os.Create(file.path)
	return fd, err
}

// copyToDir copies the source files to the destination directory.
// It recurs if -R is set and a source is a subdirectory.
func (s *State) copyToDir(cs *copyState, src []cpFile, dir cpFile) {
	for _, from := range src {
		dstPath := path.Join(upspin.PathName(dir.path), filepath.Base(from.path))
		if dir.isUpspin && from.isUpspin {
			// Try a fast copy. It can fail but that's OK.
			cs.logf("try fast copy to %s", dstPath)
			if s.fastCopy(upspin.PathName(from.path), dstPath) == nil {
				continue
			}
		}
		reader, err := s.open(from)
		if cs.recur && errors.Match(errIsDir, err) {
			// If the problem is that from is a directory but we have -R,
			// recur on the contents.
			cs.logf("recursive descent into %s", from.path)
			newFiles, err := s.contents(cs, from)
			if len(newFiles) == 0 && err != nil {
				continue
			}
			// May need to make subdirectory (even if it will have no files).
			subDir := dir
			if dir.isUpspin {
				// Rather than use the libraries and a lot of casting, it's easiest just to cat the strings here.
				subDir.path = subDir.path + "/" + filepath.Base(from.path) // TODO: is filepath.Base OK?
				_, err := s.client.MakeDirectory(upspin.PathName(subDir.path))
				if err != nil && !errors.Match(errExist, err) {
					s.fail(err)
					continue
				}
			} else {
				subDir.path = filepath.Join(subDir.path, filepath.Base(from.path))
				err := os.Mkdir(subDir.path, 0755) // TODO: Mode.
				if err != nil && !os.IsExist(err) {
					s.fail(err)
					continue
				}
			}
			s.copyToDir(cs, newFiles, subDir)
			continue
		}
		if err != nil {
			s.fail(err)
			continue
		}
		dst := cpFile{
			path:     string(dstPath),
			isUpspin: dir.isUpspin,
		}
		s.copyToFile(cs, reader, from, dst)
	}
}

// copyToFile copies the source to the destination. The source file has already been opened.
func (s *State) copyToFile(cs *copyState, reader io.ReadCloser, src, dst cpFile) {
	cs.logf("start cp %s %s", src.path, dst.path)
	defer cs.logf("end cp %s %s", src.path, dst.path)
	// If both are in Upspin, we can avoid touching the data by copying
	// just the references.
	if src.isUpspin && dst.isUpspin {
		cs.logf("try fast copy to %v", dst)
		err := s.fastCopy(upspin.PathName(src.path), upspin.PathName(dst.path))
		if err == nil {
			return
		}
	}
	writer, err := s.create(dst)
	if err != nil {
		s.fail(err)
		reader.Close()
		return
	}
	cs.doCopy(reader, writer)
}

// fastCopy copies the source to the destination using the references rather than the data.
// If it fails, PutDuplicate failed because the file exists or the source is a directory.
// (Any other error is unexpected and exits the copy command.)
// The caller may be able to retry with a regular copy.
func (s *State) fastCopy(src, dst upspin.PathName) error {
	_, err := s.client.PutDuplicate(src, dst)
	if err == nil {
		return nil
	}
	if errors.Match(errExist, err) {
		// File already exists, which PutDuplicate doesn't handle.
		// Use regular copy. We could remove it and retry
		// but that's a little scary.
		return err
	}
	if errors.Match(errIsDir, err) {
		// Oops, we have a directory. Retry.
		return err
	}
	// Unexpected error. Die.
	s.fail(err)
	return nil
}

func (cs *copyState) doCopy(reader io.ReadCloser, writer io.WriteCloser) {
	defer func() {
		reader.Close()
		err := writer.Close()
		if err != nil {
			cs.state.fail(err)
		}
	}()
	_, err := io.Copy(writer, reader)
	if err != nil {
		cs.state.fail(err)
	}
}

// glob glob-expands the argument, which could be a local file
// name or an Upspin path name.
func (cs *copyState) glob(pattern string) (files []cpFile) {
	parsed, err := path.Parse(upspin.PathName(pattern))
	if err == nil {
		// It's an Upspin path.
		for _, path := range cs.state.globUpspinPath(parsed.String()) {
			files = append(files, cpFile{
				path:     string(path),
				isUpspin: true,
			})
		}
		return files
	}
	// It's a local path.
	for _, path := range cs.state.globLocal(pattern) {
		files = append(files, cpFile{
			path:     path,
			isUpspin: false,
		})
	}
	return files
}

// contents return the top-level contents of dir as a slice of cpFiles.
func (s *State) contents(cs *copyState, dir cpFile) ([]cpFile, error) {
	if dir.isUpspin {
		entries, err := s.client.Glob(upspin.AllFilesGlob(upspin.PathName(dir.path)))
		if err != nil {
			s.fail(err)
			// OK to continue; there may still be files.
		}
		files := make([]cpFile, len(entries))
		for i, entry := range entries {
			files[i] = cpFile{
				path:     string(entry.Name),
				isUpspin: true,
			}
		}
		return files, err
	}
	// Local directory.
	fd, err := os.Open(dir.path)
	if err != nil {
		s.fail(err)
		return nil, err
	}
	defer fd.Close()
	names, err := fd.Readdirnames(0)
	if err != nil {
		s.fail(err)
		// OK to continue; there may still be files.
	}
	files := make([]cpFile, len(names))
	for i, name := range names {
		files[i] = cpFile{
			path:     filepath.Join(dir.path, name),
			isUpspin: false,
		}
	}
	return files, err
}
<commit_msg>cmd/upspin: fix nil crash in cp<commit_after>// Copyright 2016 The Upspin 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 (
	"flag"
	"io"
	"log"
	"os"
	"path/filepath"

	"upspin.io/errors"
	"upspin.io/path"
	"upspin.io/upspin"
)

func (s *State) cp(args ...string) {
	const help = `
Cp copies files into, out of, and within Upspin. If the final
argument is a directory, the files are placed inside it.  The other
arguments must not be directories unless the -R flag is set.

If the final argument is not a directory, cp requires exactly two
path names and copies the contents of the first to the second.
The -R flag requires that the final argument be a directory.

When copying from one Upspin path to another Upspin path, cp can be
very efficient, copying only the references to the data rather than
the data itself.
`
	fs := flag.NewFlagSet("cp", flag.ExitOnError)
	fs.Bool("v", false, "log each file as it is copied")
	fs.Bool("R", false, "recursively copy directories")
	s.parseFlags(fs, args, help, "cp [opts] file... file or cp [opts] file... directory")

	cs := &copyState{
		state:   s,
		flagSet: fs,
		recur:   boolFlag(fs, "R"),
		verbose: boolFlag(fs, "v"),
	}

	// Do all the glob processing here.
	// Special one-at-time glob processing because each item may be local or Upspin.
	var files []cpFile
	for _, file := range fs.Args() {
		files = append(files, cs.glob(file)...)
	}

	if len(files) < 2 {
		fs.Usage()
	}

	nSrc := len(files) - 1
	src, dest := files[:nSrc], files[nSrc]
	s.copyCommand(cs, src, dest)
}

type copyState struct {
	state   *State
	flagSet *flag.FlagSet // Used only to call Usage.
	verbose bool
	recur   bool
}

func (c *copyState) logf(format string, args ...interface{}) {
	if c.verbose {
		log.Printf(format, args...)
	}
}

// A cpFile is a glob-expanded file name and an indication of whether
// it resides on Upspin.
type cpFile struct {
	path     string
	isUpspin bool
}

var (
	errExist    = errors.E(errors.Exist)
	errNotExist = errors.E(errors.NotExist)
	errIsDir    = errors.E(errors.IsDir)
)

func (s *State) copyCommand(cs *copyState, srcFiles []cpFile, dstFile cpFile) {
	// TODO: Check for nugatory copies.
	if s.isDir(dstFile) {
		s.copyToDir(cs, srcFiles, dstFile)
		return
	}
	if len(srcFiles) != 1 {
		s.failf("copying multiple files but %s is not a directory", dstFile.path)
		cs.flagSet.Usage()
	}
	if cs.recur {
		s.failf("recursive copy requires that final argument (%s) be an existing directory", dstFile.path)
		cs.flagSet.Usage()
	}
	reader, err := s.open(srcFiles[0])
	if err != nil {
		s.exit(err)
	}
	s.copyToFile(cs, reader, srcFiles[0], dstFile)
}

// isDir reports whether the file is a directory either in Upspin
// or in the local file system.
func (s *State) isDir(cf cpFile) bool {
	if cf.isUpspin {
		entry, err := s.client.Lookup(upspin.PathName(cf.path), true)
		// Report the error here if it's anything odd, because otherwise
		// we'll report "not a directory" misleadingly.
		if err != nil && !errors.Match(errNotExist, err) {
			log.Printf("%q: %v", cf.path, err)
		}
		return err == nil && entry.IsDir()
	}
	// Not an Upspin name. Is it a local directory?
	info, err := os.Stat(cf.path)
	return err == nil && info.IsDir()
}

// open opens the file regardless of its location.
func (s *State) open(file cpFile) (io.ReadCloser, error) {
	if s.isDir(file) {
		return nil, errors.E(upspin.PathName(file.path), errors.IsDir)
	}
	if file.isUpspin {
		return s.client.Open(upspin.PathName(file.path))
	}
	return os.Open(file.path)
}

// create creates the file regardless of its location.
func (s *State) create(file cpFile) (io.WriteCloser, error) {
	if file.isUpspin {
		fd, err := s.client.Create(upspin.PathName(file.path))
		return fd, err
	}
	fd, err := os.Create(file.path)
	return fd, err
}

// copyToDir copies the source files to the destination directory.
// It recurs if -R is set and a source is a subdirectory.
func (s *State) copyToDir(cs *copyState, src []cpFile, dir cpFile) {
	for _, from := range src {
		dstPath := path.Join(upspin.PathName(dir.path), filepath.Base(from.path))
		if dir.isUpspin && from.isUpspin {
			// Try a fast copy. It can fail but that's OK.
			cs.logf("try fast copy to %s", dstPath)
			if s.fastCopy(upspin.PathName(from.path), dstPath) == nil {
				continue
			}
		}
		reader, err := s.open(from)
		if cs.recur && errors.Match(errIsDir, err) {
			// If the problem is that from is a directory but we have -R,
			// recur on the contents.
			cs.logf("recursive descent into %s", from.path)
			newFiles, err := s.contents(cs, from)
			if len(newFiles) == 0 && err != nil {
				continue
			}
			// May need to make subdirectory (even if it will have no files).
			subDir := dir
			if dir.isUpspin {
				// Rather than use the libraries and a lot of casting, it's easiest just to cat the strings here.
				subDir.path = subDir.path + "/" + filepath.Base(from.path) // TODO: is filepath.Base OK?
				_, err := s.client.MakeDirectory(upspin.PathName(subDir.path))
				if err != nil && !errors.Match(errExist, err) {
					s.fail(err)
					continue
				}
			} else {
				subDir.path = filepath.Join(subDir.path, filepath.Base(from.path))
				err := os.Mkdir(subDir.path, 0755) // TODO: Mode.
				if err != nil && !os.IsExist(err) {
					s.fail(err)
					continue
				}
			}
			s.copyToDir(cs, newFiles, subDir)
			continue
		}
		if err != nil {
			s.fail(err)
			continue
		}
		dst := cpFile{
			path:     string(dstPath),
			isUpspin: dir.isUpspin,
		}
		s.copyToFile(cs, reader, from, dst)
	}
}

// copyToFile copies the source to the destination. The source file has already been opened.
func (s *State) copyToFile(cs *copyState, reader io.ReadCloser, src, dst cpFile) {
	cs.logf("start cp %s %s", src.path, dst.path)
	defer cs.logf("end cp %s %s", src.path, dst.path)
	// If both are in Upspin, we can avoid touching the data by copying
	// just the references.
	if src.isUpspin && dst.isUpspin {
		cs.logf("try fast copy to %v", dst)
		err := s.fastCopy(upspin.PathName(src.path), upspin.PathName(dst.path))
		if err == nil {
			return
		}
	}
	writer, err := s.create(dst)
	if err != nil {
		s.fail(err)
		reader.Close()
		return
	}
	cs.doCopy(reader, writer)
}

// fastCopy copies the source to the destination using the references rather than the data.
// If it fails, PutDuplicate failed because the file exists or the source is a directory.
// (Any other error is unexpected and exits the copy command.)
// The caller may be able to retry with a regular copy.
func (s *State) fastCopy(src, dst upspin.PathName) error {
	_, err := s.client.PutDuplicate(src, dst)
	if err == nil {
		return nil
	}
	if errors.Match(errExist, err) {
		// File already exists, which PutDuplicate doesn't handle.
		// Use regular copy. We could remove it and retry
		// but that's a little scary.
		return err
	}
	if errors.Match(errIsDir, err) {
		// Oops, we have a directory. Retry.
		return err
	}
	// Unexpected error. Die.
	s.fail(err)
	return nil
}

func (cs *copyState) doCopy(reader io.ReadCloser, writer io.WriteCloser) {
	defer func() {
		reader.Close()
		err := writer.Close()
		if err != nil {
			cs.state.fail(err)
		}
	}()
	_, err := io.Copy(writer, reader)
	if err != nil {
		cs.state.fail(err)
	}
}

// glob glob-expands the argument, which could be a local file
// name or an Upspin path name.
func (cs *copyState) glob(pattern string) (files []cpFile) {
	parsed, err := path.Parse(upspin.PathName(pattern))
	if err == nil {
		// It's an Upspin path.
		for _, path := range cs.state.globUpspinPath(parsed.String()) {
			files = append(files, cpFile{
				path:     string(path),
				isUpspin: true,
			})
		}
		return files
	}
	// It's a local path.
	for _, path := range cs.state.globLocal(pattern) {
		files = append(files, cpFile{
			path:     path,
			isUpspin: false,
		})
	}
	return files
}

// contents return the top-level contents of dir as a slice of cpFiles.
func (s *State) contents(cs *copyState, dir cpFile) ([]cpFile, error) {
	if dir.isUpspin {
		entries, err := s.client.Glob(upspin.AllFilesGlob(upspin.PathName(dir.path)))
		if err != nil {
			s.fail(err)
			// OK to continue; there may still be files.
		}
		files := make([]cpFile, len(entries))
		for i, entry := range entries {
			files[i] = cpFile{
				path:     string(entry.Name),
				isUpspin: true,
			}
		}
		return files, err
	}
	// Local directory.
	fd, err := os.Open(dir.path)
	if err != nil {
		s.fail(err)
		return nil, err
	}
	defer fd.Close()
	names, err := fd.Readdirnames(0)
	if err != nil {
		s.fail(err)
		// OK to continue; there may still be files.
	}
	files := make([]cpFile, len(names))
	for i, name := range names {
		files[i] = cpFile{
			path:     filepath.Join(dir.path, name),
			isUpspin: false,
		}
	}
	return files, err
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	_ "image/jpeg"
	"image/png"
	"io/ioutil"
	"os"
)

func main() {
	data, err := ioutil.ReadFile("input.jpg")

	if err != nil {
		panic(err)
	}

	img, _, err := image.Decode(bytes.NewReader(data))

	if err != nil {
		panic(err)
	}

	improved := ImproveQuality(img)
	f, err := os.Create("output.png")

	if err != nil {
		panic(err)
	}

	defer f.Close()
	png.Encode(f, improved)
}

// Pixel ...
type Pixel struct {
	X int
	Y int
}

// Area ...
type Area struct {
	Pixels []Pixel
	totalR uint64
	totalG uint64
	totalB uint64
	totalA uint64
}

// Add ...
func (area *Area) Add(x, y int, r, g, b, a uint32) {
	area.Pixels = append(area.Pixels, Pixel{
		X: x,
		Y: y,
	})

	area.totalR += uint64(r)
	area.totalG += uint64(g)
	area.totalB += uint64(b)
	area.totalA += uint64(a)
}

// AverageColor ...
func (area *Area) AverageColor() color.Color {
	if len(area.Pixels) == 0 {
		return color.Transparent
	}

	return color.RGBA64{
		R: uint16(area.totalR / uint64(len(area.Pixels))),
		G: uint16(area.totalG / uint64(len(area.Pixels))),
		B: uint16(area.totalB / uint64(len(area.Pixels))),
		A: uint16(area.totalA / uint64(len(area.Pixels))),
	}
}

const (
	tolerance = uint32(3000)
)

func diffAbs(a uint32, b uint32) uint32 {
	if a > b {
		return a - b
	}

	return b - a
}

// ImproveQuality returns the average color of an image in HSL format.
func ImproveQuality(img image.Image) *image.NRGBA {
	width := img.Bounds().Dx()
	height := img.Bounds().Dy()
	clone := image.NewNRGBA(image.Rect(0, 0, width, height))
	areas := []*Area{}
	areaIndexMap := make([]int, width*height)

	for x := 0; x < width; x++ {
		for y := 0; y < height; y++ {
			color := img.At(x, y)
			r, g, b, a := color.RGBA()
			areaIndex := -1

			// Find similar area
			for i := 0; i < len(areas); i++ {
				area := areas[i]
				avgR, avgG, avgB, _ := area.AverageColor().RGBA()

				// Is the color similar?
				if diffAbs(r, avgR) <= tolerance && diffAbs(g, avgG) <= tolerance && diffAbs(b, avgB) <= tolerance {
					areaIndex = i
					break
				}
			}

			// Insert new area
			if areaIndex == -1 {
				areaIndex = len(areas)
				areas = append(areas, &Area{})
			}

			areaIndexMap[y*width+x] = areaIndex
			areas[areaIndex].Add(x, y, r, g, b, a)
		}
	}

	fmt.Println(len(areas), "areas")

	// Reduce noise
	noiseCount := 0

	for areaIndex, area := range areas {
		noisePixelIndices := []int{}
		areaSurroundedBy := map[int]int{}

		for i := 0; i < len(area.Pixels); i++ {
			// If pixel is surrounded by 4 different areas, remove it
			pixel := area.Pixels[i]
			x := pixel.X
			y := pixel.Y
			left := areaIndex
			right := areaIndex
			top := areaIndex
			bottom := areaIndex

			if x > 0 {
				left = areaIndexMap[y*width+(x-1)]
			}

			if x < width-1 {
				right = areaIndexMap[y*width+(x+1)]
			}

			if y > 0 {
				top = areaIndexMap[(y-1)*width+x]
			}

			if y < height-1 {
				bottom = areaIndexMap[(y+1)*width+x]
			}

			differentNeighbors := 0

			if left != areaIndex {
				differentNeighbors++
			}

			if right != areaIndex {
				differentNeighbors++
			}

			if top != areaIndex {
				differentNeighbors++
			}

			if bottom != areaIndex {
				differentNeighbors++
			}

			// Determine surrounding area
			areaIndexScore := map[int]int{}

			areaIndexScore[left]++
			areaIndexScore[right]++
			areaIndexScore[top]++
			areaIndexScore[bottom]++

			areaSurroundedBy[left]++
			areaSurroundedBy[right]++
			areaSurroundedBy[top]++
			areaSurroundedBy[bottom]++

			newAreaIndex := -1
			bestScore := 0

			for checkIndex, score := range areaIndexScore {
				if score > bestScore {
					bestScore = score
					newAreaIndex = checkIndex
				}
			}

			if differentNeighbors >= 3 && bestScore >= 3 {
				noiseCount++
				noisePixelIndices = append(noisePixelIndices, i)

				// Add to surrounding area
				r, g, b, a := img.At(x, y).RGBA()
				areas[newAreaIndex].Add(x, y, r, g, b, a)

				area.totalR -= uint64(r)
				area.totalG -= uint64(g)
				area.totalB -= uint64(b)
				area.totalA -= uint64(a)
			}
		}

		// Remove noise pixels
		offset := 0

		for _, removal := range noisePixelIndices {
			index := removal - offset
			area.Pixels = append(area.Pixels[:index], area.Pixels[index+1:]...)
			offset++
		}

		// // Determine surrounding area
		// surroundingAreaIndex := -1
		// bestScore := 0

		// for checkIndex, score := range areaSurroundedBy {
		// 	if score > bestScore && checkIndex != areaIndex {
		// 		bestScore = score
		// 		surroundingAreaIndex = checkIndex
		// 	}
		// }

		// surroundingArea := areas[surroundingAreaIndex]

		// if areaIndex != surroundingAreaIndex && len(surroundingArea.Pixels) > len(area.Pixels)*2 {
		// 	// const surroundTolerance = 5000

		// 	// r1, g1, b1, a1 := area.AverageColor().RGBA()
		// 	// r2, g2, b2, a2 := surroundingArea.AverageColor().RGBA()

		// 	// if diffAbs(r1, r2) < surroundTolerance && diffAbs(g1, g2) < surroundTolerance && diffAbs(b1, b2) < surroundTolerance && diffAbs(a1, a2) < surroundTolerance {
		// 	// 	// fmt.Println(areaIndex, "surrounded by", surroundingAreaIndex, "|", len(area.Pixels), len(surroundingArea.Pixels))

		// 	// 	// Add pixels to surrounding area
		// 	// 	for _, pixel := range area.Pixels {
		// 	// 		r, g, b, a := img.At(pixel.X, pixel.Y).RGBA()
		// 	// 		surroundingArea.Add(pixel.X, pixel.Y, r, g, b, a)
		// 	// 	}

		// 	// 	// Remove this area
		// 	// 	area.Pixels = nil
		// 	// 	area.totalR = 0
		// 	// 	area.totalG = 0
		// 	// 	area.totalB = 0
		// 	// 	area.totalA = 0
		// 	// }
		// }
	}

	fmt.Println(noiseCount, "noise pixels")

	pixelCount := 0

	for _, area := range areas {
		pixelCount += len(area.Pixels)
	}

	fmt.Println(pixelCount, "pixels", width*height)

	// Build image from areas
	for _, area := range areas {
		avgColor := area.AverageColor()

		for _, pixel := range area.Pixels {
			clone.Set(pixel.X, pixel.Y, avgColor)
		}
	}

	return clone
}
<commit_msg>Disabled image quality patch<commit_after>package main

func main() {}

// import (
// 	"bytes"
// 	"fmt"
// 	"image"
// 	"image/color"
// 	_ "image/jpeg"
// 	"image/png"
// 	"io/ioutil"
// 	"os"
// )

// func main() {
// 	data, err := ioutil.ReadFile("input.jpg")

// 	if err != nil {
// 		panic(err)
// 	}

// 	img, _, err := image.Decode(bytes.NewReader(data))

// 	if err != nil {
// 		panic(err)
// 	}

// 	improved := ImproveQuality(img)
// 	f, err := os.Create("output.png")

// 	if err != nil {
// 		panic(err)
// 	}

// 	defer f.Close()
// 	png.Encode(f, improved)
// }

// // Pixel ...
// type Pixel struct {
// 	X int
// 	Y int
// }

// // Area ...
// type Area struct {
// 	Pixels []Pixel
// 	totalR uint64
// 	totalG uint64
// 	totalB uint64
// 	totalA uint64
// }

// // Add ...
// func (area *Area) Add(x, y int, r, g, b, a uint32) {
// 	area.Pixels = append(area.Pixels, Pixel{
// 		X: x,
// 		Y: y,
// 	})

// 	area.totalR += uint64(r)
// 	area.totalG += uint64(g)
// 	area.totalB += uint64(b)
// 	area.totalA += uint64(a)
// }

// // AverageColor ...
// func (area *Area) AverageColor() color.Color {
// 	if len(area.Pixels) == 0 {
// 		return color.Transparent
// 	}

// 	return color.RGBA64{
// 		R: uint16(area.totalR / uint64(len(area.Pixels))),
// 		G: uint16(area.totalG / uint64(len(area.Pixels))),
// 		B: uint16(area.totalB / uint64(len(area.Pixels))),
// 		A: uint16(area.totalA / uint64(len(area.Pixels))),
// 	}
// }

// const (
// 	tolerance = uint32(3000)
// )

// func diffAbs(a uint32, b uint32) uint32 {
// 	if a > b {
// 		return a - b
// 	}

// 	return b - a
// }

// // ImproveQuality returns the average color of an image in HSL format.
// func ImproveQuality(img image.Image) *image.NRGBA {
// 	width := img.Bounds().Dx()
// 	height := img.Bounds().Dy()
// 	clone := image.NewNRGBA(image.Rect(0, 0, width, height))
// 	areas := []*Area{}
// 	areaIndexMap := make([]int, width*height)

// 	for x := 0; x < width; x++ {
// 		for y := 0; y < height; y++ {
// 			color := img.At(x, y)
// 			r, g, b, a := color.RGBA()
// 			areaIndex := -1

// 			// Find similar area
// 			for i := 0; i < len(areas); i++ {
// 				area := areas[i]
// 				avgR, avgG, avgB, _ := area.AverageColor().RGBA()

// 				// Is the color similar?
// 				if diffAbs(r, avgR) <= tolerance && diffAbs(g, avgG) <= tolerance && diffAbs(b, avgB) <= tolerance {
// 					areaIndex = i
// 					break
// 				}
// 			}

// 			// Insert new area
// 			if areaIndex == -1 {
// 				areaIndex = len(areas)
// 				areas = append(areas, &Area{})
// 			}

// 			areaIndexMap[y*width+x] = areaIndex
// 			areas[areaIndex].Add(x, y, r, g, b, a)
// 		}
// 	}

// 	fmt.Println(len(areas), "areas")

// 	// Reduce noise
// 	noiseCount := 0

// 	for areaIndex, area := range areas {
// 		noisePixelIndices := []int{}
// 		areaSurroundedBy := map[int]int{}

// 		for i := 0; i < len(area.Pixels); i++ {
// 			// If pixel is surrounded by 4 different areas, remove it
// 			pixel := area.Pixels[i]
// 			x := pixel.X
// 			y := pixel.Y
// 			left := areaIndex
// 			right := areaIndex
// 			top := areaIndex
// 			bottom := areaIndex

// 			if x > 0 {
// 				left = areaIndexMap[y*width+(x-1)]
// 			}

// 			if x < width-1 {
// 				right = areaIndexMap[y*width+(x+1)]
// 			}

// 			if y > 0 {
// 				top = areaIndexMap[(y-1)*width+x]
// 			}

// 			if y < height-1 {
// 				bottom = areaIndexMap[(y+1)*width+x]
// 			}

// 			differentNeighbors := 0

// 			if left != areaIndex {
// 				differentNeighbors++
// 			}

// 			if right != areaIndex {
// 				differentNeighbors++
// 			}

// 			if top != areaIndex {
// 				differentNeighbors++
// 			}

// 			if bottom != areaIndex {
// 				differentNeighbors++
// 			}

// 			// Determine surrounding area
// 			areaIndexScore := map[int]int{}

// 			areaIndexScore[left]++
// 			areaIndexScore[right]++
// 			areaIndexScore[top]++
// 			areaIndexScore[bottom]++

// 			areaSurroundedBy[left]++
// 			areaSurroundedBy[right]++
// 			areaSurroundedBy[top]++
// 			areaSurroundedBy[bottom]++

// 			newAreaIndex := -1
// 			bestScore := 0

// 			for checkIndex, score := range areaIndexScore {
// 				if score > bestScore {
// 					bestScore = score
// 					newAreaIndex = checkIndex
// 				}
// 			}

// 			if differentNeighbors >= 3 && bestScore >= 3 {
// 				noiseCount++
// 				noisePixelIndices = append(noisePixelIndices, i)

// 				// Add to surrounding area
// 				r, g, b, a := img.At(x, y).RGBA()
// 				areas[newAreaIndex].Add(x, y, r, g, b, a)

// 				area.totalR -= uint64(r)
// 				area.totalG -= uint64(g)
// 				area.totalB -= uint64(b)
// 				area.totalA -= uint64(a)
// 			}
// 		}

// 		// Remove noise pixels
// 		offset := 0

// 		for _, removal := range noisePixelIndices {
// 			index := removal - offset
// 			area.Pixels = append(area.Pixels[:index], area.Pixels[index+1:]...)
// 			offset++
// 		}

// 		// // Determine surrounding area
// 		// surroundingAreaIndex := -1
// 		// bestScore := 0

// 		// for checkIndex, score := range areaSurroundedBy {
// 		// 	if score > bestScore && checkIndex != areaIndex {
// 		// 		bestScore = score
// 		// 		surroundingAreaIndex = checkIndex
// 		// 	}
// 		// }

// 		// surroundingArea := areas[surroundingAreaIndex]

// 		// if areaIndex != surroundingAreaIndex && len(surroundingArea.Pixels) > len(area.Pixels)*2 {
// 		// 	// const surroundTolerance = 5000

// 		// 	// r1, g1, b1, a1 := area.AverageColor().RGBA()
// 		// 	// r2, g2, b2, a2 := surroundingArea.AverageColor().RGBA()

// 		// 	// if diffAbs(r1, r2) < surroundTolerance && diffAbs(g1, g2) < surroundTolerance && diffAbs(b1, b2) < surroundTolerance && diffAbs(a1, a2) < surroundTolerance {
// 		// 	// 	// fmt.Println(areaIndex, "surrounded by", surroundingAreaIndex, "|", len(area.Pixels), len(surroundingArea.Pixels))

// 		// 	// 	// Add pixels to surrounding area
// 		// 	// 	for _, pixel := range area.Pixels {
// 		// 	// 		r, g, b, a := img.At(pixel.X, pixel.Y).RGBA()
// 		// 	// 		surroundingArea.Add(pixel.X, pixel.Y, r, g, b, a)
// 		// 	// 	}

// 		// 	// 	// Remove this area
// 		// 	// 	area.Pixels = nil
// 		// 	// 	area.totalR = 0
// 		// 	// 	area.totalG = 0
// 		// 	// 	area.totalB = 0
// 		// 	// 	area.totalA = 0
// 		// 	// }
// 		// }
// 	}

// 	fmt.Println(noiseCount, "noise pixels")

// 	pixelCount := 0

// 	for _, area := range areas {
// 		pixelCount += len(area.Pixels)
// 	}

// 	fmt.Println(pixelCount, "pixels", width*height)

// 	// Build image from areas
// 	for _, area := range areas {
// 		avgColor := area.AverageColor()

// 		for _, pixel := range area.Pixels {
// 			clone.Set(pixel.X, pixel.Y, avgColor)
// 		}
// 	}

// 	return clone
// }
<|endoftext|>


<commit_before>// Copyright 2022 The etcd 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 mvcc

import (
	"hash"
	"hash/crc32"

	"go.etcd.io/etcd/server/v3/storage/backend"
	"go.etcd.io/etcd/server/v3/storage/schema"
)

func unsafeHashByRev(tx backend.ReadTx, lower, upper revision, keep map[revision]struct{}) (uint32, error) {
	h := newKVHasher()
	err := tx.UnsafeForEach(schema.Key, func(k, v []byte) error {
		kr := bytesToRev(k)
		if !upper.GreaterThan(kr) {
			return nil
		}
		// skip revisions that are scheduled for deletion
		// due to compacting; don't skip if there isn't one.
		if lower.GreaterThan(kr) && len(keep) > 0 {
			if _, ok := keep[kr]; !ok {
				return nil
			}
		}
		h.WriteKeyValue(k, v)
		return nil
	})
	return h.Hash(), err
}

type hasher struct {
	h hash.Hash32
}

func newKVHasher() hasher {
	h := crc32.New(crc32.MakeTable(crc32.Castagnoli))
	h.Write(schema.Key.Name())
	return hasher{h}
}

func (h *hasher) WriteKeyValue(k, v []byte) {
	h.h.Write(k)
	h.h.Write(v)
}

func (h *hasher) Hash() uint32 {
	return h.h.Sum32()
}
<commit_msg>server: Refactor hasher<commit_after>// Copyright 2022 The etcd 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 mvcc

import (
	"hash"
	"hash/crc32"

	"go.etcd.io/etcd/server/v3/storage/backend"
	"go.etcd.io/etcd/server/v3/storage/schema"
)

func unsafeHashByRev(tx backend.ReadTx, lower, upper revision, keep map[revision]struct{}) (uint32, error) {
	h := newKVHasher(lower, upper, keep)
	err := tx.UnsafeForEach(schema.Key, func(k, v []byte) error {
		h.WriteKeyValue(k, v)
		return nil
	})
	return h.Hash(), err
}

type kvHasher struct {
	hash         hash.Hash32
	lower, upper revision
	keep         map[revision]struct{}
}

func newKVHasher(lower, upper revision, keep map[revision]struct{}) kvHasher {
	h := crc32.New(crc32.MakeTable(crc32.Castagnoli))
	h.Write(schema.Key.Name())
	return kvHasher{
		hash:  h,
		lower: lower,
		upper: upper,
		keep:  keep,
	}
}

func (h *kvHasher) WriteKeyValue(k, v []byte) {
	kr := bytesToRev(k)
	if !h.upper.GreaterThan(kr) {
		return
	}
	// skip revisions that are scheduled for deletion
	// due to compacting; don't skip if there isn't one.
	if h.lower.GreaterThan(kr) && len(h.keep) > 0 {
		if _, ok := h.keep[kr]; !ok {
			return
		}
	}
	h.hash.Write(k)
	h.hash.Write(v)
}

func (h *kvHasher) Hash() uint32 {
	return h.hash.Sum32()
}
<|endoftext|>


<commit_before>// Package coffeebeanbot is a coffee bean inspired bot created to help me through my day.
// Its current focus is to handle "Pomodoro Technique"-style timeboxing notification.
package coffeebeanbot

// Copyright 2017 Sean A. Pfeifer

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/bwmarrin/discordgo"
)

const (
	discordBotPrefix    = "Bot "
	pomDuration         = time.Minute * 25
	voiceWaitTime       = time.Millisecond * 250 // The amount of time to sleep before speaking & leaving the voice channel
	baseAuthURLTemplate = "https://discordapp.com/api/oauth2/authorize?client_id=%s&scope=bot"
)

// cmdHandler is the type for our functions that will be called upon receiving commands from a user.
type cmdHandler func(s *discordgo.Session, m *discordgo.MessageCreate, extra string)

type botCommand struct {
	handler       cmdHandler
	desc          string
	exampleParams string
}

// Bot contains the information needed to run the Discord bot
type Bot struct {
	Config      Config
	started     time.Time
	cmdHandlers map[string]botCommand
	discord     *discordgo.Session

	helpMessage        string
	inviteMessage      string
	poms               channelPomMap
	workEndAudioBuffer [][]byte
}

// Config is the Bot's configuration data
type Config struct {
	AuthToken    string `json:"authToken"` // AuthToken is all that we need to authenticate as the bot
	ClientID     string `json:"clientID"`  // ClientID is used to create the invite link for the bot, and isn't necessary for login
	CmdPrefix    string `json:"cmdPrefix"`
	WorkEndAudio string `json:"workEndAudio"`
}

// LoadConfigFile loads the config from the given path, returning the config or an error if one occurred.
// I generally prefer config files over environment variables, due to the ease of setting them up as secrets
// in Kubernetes.
func LoadConfigFile(path string) (*Config, error) {
	data, err := ioutil.ReadFile(path)

	if err != nil {
		return nil, err
	}

	var cfg Config
	err = json.Unmarshal(data, &cfg)

	return &cfg, err
}

// NewBot is how you should create a new Bot in order to assure that all initialization has been completed.
func NewBot(config Config) *Bot {
	bot := &Bot{
		Config: config,
		poms:   newChannelPomMap(),
	}

	bot.registerCmdHandlers()
	bot.inviteMessage = fmt.Sprintf("To have me join your server, click here: <"+baseAuthURLTemplate+">", bot.Config.ClientID)
	bot.helpMessage = bot.buildHelpMessage()
	bot.loadSounds()

	return bot
}

func (bot *Bot) loadSounds() {
	audioBuffer, err := LoadDiscordAudio(bot.Config.WorkEndAudio)
	if err != nil {
		log.Printf("Error loading audio: %v", err)
	} else {
		bot.workEndAudioBuffer = audioBuffer
	}
}

func (bot *Bot) registerCmdHandlers() {
	bot.cmdHandlers = map[string]botCommand{
		"ping":    {handler: bot.onCmdPing, desc: "Pings the bot to ensure it is currently running", exampleParams: ""},
		"started": {handler: bot.onCmdStarted, desc: "Shows when the current version of the bot started running", exampleParams: ""},
		"invite":  {handler: bot.onCmdInvite, desc: "Creates an invite link you can use to have the bot join your server", exampleParams: ""},
		"echo":    {handler: bot.onCmdEcho, desc: "Echoes the given message back", exampleParams: "Hello, world!"},
		"start":   {handler: bot.onCmdStartPom, desc: "Starts a Pomodoro work cycle on the channel", exampleParams: ""},
		"cancel":  {handler: bot.onCmdCancelPom, desc: "Cancels the current Pomodoro work cycle on the channel", exampleParams: ""},
		"help":    {handler: bot.onCmdHelp, desc: "Shows this help message", exampleParams: ""},
	}
}

func (bot *Bot) buildHelpMessage() string {
	helpBuf := bytes.Buffer{}
	helpBuf.WriteString("This bot was written by Sean A. Pfeifer to help him get more done.\n")

	// I don't really care about ordering right now - this is intentionally using the map iteration order,
	// which I am aware is pseudo-random.
	// TODO: Add a "group" attribute to the commands, and sort by group, then command.
	for cmdStr, cmd := range bot.cmdHandlers {
		helpBuf.WriteString(fmt.Sprintf("\n•  **%s**  -  %s\n", cmdStr, cmd.desc))
		helpBuf.WriteString(fmt.Sprintf("    Example: `%s%s %s`\n", bot.Config.CmdPrefix, cmdStr, cmd.exampleParams))
	}

	helpBuf.WriteString("\n" + bot.inviteMessage)

	return helpBuf.String()
}

// Start will start the bot, blocking until completion
func (bot *Bot) Start() error {
	if bot.Config.AuthToken == "" {
		return errors.New("no auth token found in config")
	}

	var err error
	bot.discord, err = discordgo.New(discordBotPrefix + bot.Config.AuthToken)
	if err != nil {
		return err
	}

	bot.discord.AddHandler(bot.onReady)
	bot.discord.AddHandler(bot.onMessageReceived)

	err = bot.discord.Open()
	if err != nil {
		return err
	}

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
	<-sc

	return bot.discord.Close()
}

func (bot *Bot) onReady(s *discordgo.Session, event *discordgo.Ready) {
	bot.started = time.Now()
	log.Printf("Bot connected and ready as '%s#%s'!", event.User.Username, event.User.Discriminator)
}

// onMessageReceived is called when a message is received on a channel that the bot is listening on.
// It will dispatch known commands to the command handlers, passing along any extra string information.
func (bot *Bot) onMessageReceived(s *discordgo.Session, m *discordgo.MessageCreate) {
	// Ignore messages created by this bot
	if m.Author.ID == s.State.User.ID {
		return
	}

	msg := m.Content

	cmdPrefixLen := len(bot.Config.CmdPrefix)

	// Dispatch the command iff we have our prefix (case-insensitive).
	if len(msg) > cmdPrefixLen && strings.EqualFold(bot.Config.CmdPrefix, msg[0:cmdPrefixLen]) {
		afterPrefix := msg[cmdPrefixLen:]
		cmd := strings.SplitN(afterPrefix, " ", 2)

		if f, ok := bot.cmdHandlers[strings.ToLower(cmd[0])]; ok {
			rest := ""
			if len(cmd) > 1 {
				rest = cmd[1]
			}

			if f.handler != nil {
				f.handler(s, m, rest)
			} else {
				log.Printf("Error: nil handler for command '%s'", cmd)
				s.ChannelMessageSend(m.ChannelID, "Command error - please contact support.")
			}
		}
	}
}

func (bot *Bot) onCmdPing(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	s.ChannelMessageSend(m.ChannelID, "Pong!")
}

func (bot *Bot) onCmdStarted(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	s.ChannelMessageSend(m.ChannelID, "I started "+bot.started.String())
}

func (bot *Bot) onCmdEcho(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	// Make sure the echoed text can't break out of our quote box.
	extra = strings.Replace(extra, "`", "", -1)
	s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("Echo:  `%s`", extra))
}

func (bot *Bot) onCmdHelp(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	s.ChannelMessageSend(m.ChannelID, bot.helpMessage)
}

func (bot *Bot) onCmdInvite(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	s.ChannelMessageSend(m.ChannelID, bot.inviteMessage)
}

func (bot *Bot) onCmdStartPom(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	channel, err := s.State.Channel(m.ChannelID)
	if err != nil {
		// Could not find the channel, so simply log and exit
		log.Printf("Could not find channel for ChannelID '%s'", m.ChannelID)
		return
	}

	notif := NotifyInfo{
		extra,
		m.Author.ID,
		channel.GuildID,
	}

	if bot.poms.CreateIfEmpty(m.ChannelID, pomDuration, func() { bot.onPomEnded(m.ChannelID) }, notif) {
		s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("Pomodoro started - **%.1f minutes** remaining!", pomDuration.Minutes()))
	} else {
		s.ChannelMessageSend(m.ChannelID, "A Pomodoro is already running on this channel.")
	}
}

func (bot *Bot) onCmdCancelPom(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	if notif := bot.poms.RemoveIfExists(m.ChannelID); notif != nil {
		// TODO: Use the NotifyInfo here?
		s.ChannelMessageSend(m.ChannelID, "Pomodoro cancelled!")
	} else {
		s.ChannelMessageSend(m.ChannelID, "No Pomodoro running on this channel to cancel.")
	}
}

// onPomEnded performs the notification
func (bot *Bot) onPomEnded(channelID string) {
	notif := bot.poms.RemoveIfExists(channelID)
	message := "Pomodoro ended.  Time for a short break!"

	var toMention []string

	if notif != nil {
		user, err := bot.discord.User(notif.UserID)
		if err == nil {
			toMention = append(toMention, user.Mention())
		}
		// Doing this in a goroutine so we don't wait until the audio has been played to send the text notification.
		// This isn't required, but is my preference.
		go bot.playEndSound(*notif)
	} else {
		bot.discord.ChannelMessageSend(channelID, message)
	}

	if len(toMention) > 0 {
		mentions := strings.Join(toMention, " ")
		bot.discord.ChannelMessageSend(channelID, fmt.Sprintf("%s %s", message, mentions))
	} else {
		bot.discord.ChannelMessageSend(channelID, message)
	}
}
<commit_msg>Removed "ping", "start", and "echo" commands, as they are not really useful. Simplified and fixed onPomEnded a bit to only send the message at one point (and only when we should). onPomEnded now sends the task title, if one if given.<commit_after>// Package coffeebeanbot is a coffee bean inspired bot created to help me through my day.
// Its current focus is to handle "Pomodoro Technique"-style timeboxing notification.
package coffeebeanbot

// Copyright 2017 Sean A. Pfeifer

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/bwmarrin/discordgo"
)

const (
	discordBotPrefix    = "Bot "
	pomDuration         = time.Minute * 25
	voiceWaitTime       = time.Millisecond * 250 // The amount of time to sleep before speaking & leaving the voice channel
	baseAuthURLTemplate = "https://discordapp.com/api/oauth2/authorize?client_id=%s&scope=bot"
)

// cmdHandler is the type for our functions that will be called upon receiving commands from a user.
type cmdHandler func(s *discordgo.Session, m *discordgo.MessageCreate, extra string)

type botCommand struct {
	handler       cmdHandler
	desc          string
	exampleParams string
}

// Bot contains the information needed to run the Discord bot
type Bot struct {
	Config      Config
	cmdHandlers map[string]botCommand
	discord     *discordgo.Session

	helpMessage        string
	inviteMessage      string
	poms               channelPomMap
	workEndAudioBuffer [][]byte
}

// Config is the Bot's configuration data
type Config struct {
	AuthToken    string `json:"authToken"` // AuthToken is all that we need to authenticate as the bot
	ClientID     string `json:"clientID"`  // ClientID is used to create the invite link for the bot, and isn't necessary for login
	CmdPrefix    string `json:"cmdPrefix"`
	WorkEndAudio string `json:"workEndAudio"`
}

// LoadConfigFile loads the config from the given path, returning the config or an error if one occurred.
// I generally prefer config files over environment variables, due to the ease of setting them up as secrets
// in Kubernetes.
func LoadConfigFile(path string) (*Config, error) {
	data, err := ioutil.ReadFile(path)

	if err != nil {
		return nil, err
	}

	var cfg Config
	err = json.Unmarshal(data, &cfg)

	return &cfg, err
}

// NewBot is how you should create a new Bot in order to assure that all initialization has been completed.
func NewBot(config Config) *Bot {
	bot := &Bot{
		Config: config,
		poms:   newChannelPomMap(),
	}

	bot.registerCmdHandlers()
	bot.inviteMessage = fmt.Sprintf("To have me join your server, click here: <"+baseAuthURLTemplate+">", bot.Config.ClientID)
	bot.helpMessage = bot.buildHelpMessage()
	bot.loadSounds()

	return bot
}

func (bot *Bot) loadSounds() {
	audioBuffer, err := LoadDiscordAudio(bot.Config.WorkEndAudio)
	if err != nil {
		log.Printf("Error loading audio: %v", err)
	} else {
		bot.workEndAudioBuffer = audioBuffer
	}
}

func (bot *Bot) registerCmdHandlers() {
	bot.cmdHandlers = map[string]botCommand{
		"invite": {handler: bot.onCmdInvite, desc: "Creates an invite link you can use to have the bot join your server", exampleParams: ""},
		"start":  {handler: bot.onCmdStartPom, desc: "Starts a Pomodoro work cycle on the channel. You can optionally specify the task you are working on", exampleParams: "Create a new notification sound, add an example"},
		"cancel": {handler: bot.onCmdCancelPom, desc: "Cancels the current Pomodoro work cycle on the channel", exampleParams: ""},
		"help":   {handler: bot.onCmdHelp, desc: "Shows this help message", exampleParams: ""},
	}
}

func (bot *Bot) buildHelpMessage() string {
	helpBuf := bytes.Buffer{}
	helpBuf.WriteString("This bot was written by Sean A. Pfeifer to help him get more done.\n")

	// I don't really care about ordering right now - this is intentionally using the map iteration order,
	// which I am aware is pseudo-random.
	// TODO: Add a "group" attribute to the commands, and sort by group, then command.
	for cmdStr, cmd := range bot.cmdHandlers {
		helpBuf.WriteString(fmt.Sprintf("\n•  **%s**  -  %s\n", cmdStr, cmd.desc))
		helpBuf.WriteString(fmt.Sprintf("    Example: `%s%s %s`\n", bot.Config.CmdPrefix, cmdStr, cmd.exampleParams))
	}

	helpBuf.WriteString("\n" + bot.inviteMessage)

	return helpBuf.String()
}

// Start will start the bot, blocking until completion
func (bot *Bot) Start() error {
	if bot.Config.AuthToken == "" {
		return errors.New("no auth token found in config")
	}

	var err error
	bot.discord, err = discordgo.New(discordBotPrefix + bot.Config.AuthToken)
	if err != nil {
		return err
	}

	bot.discord.AddHandler(bot.onReady)
	bot.discord.AddHandler(bot.onMessageReceived)

	err = bot.discord.Open()
	if err != nil {
		return err
	}

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
	<-sc

	return bot.discord.Close()
}

func (bot *Bot) onReady(s *discordgo.Session, event *discordgo.Ready) {
	log.Printf("Bot connected and ready as '%s#%s'!", event.User.Username, event.User.Discriminator)
}

// onMessageReceived is called when a message is received on a channel that the bot is listening on.
// It will dispatch known commands to the command handlers, passing along any extra string information.
func (bot *Bot) onMessageReceived(s *discordgo.Session, m *discordgo.MessageCreate) {
	// Ignore messages created by this bot
	if m.Author.ID == s.State.User.ID {
		return
	}

	msg := m.Content

	cmdPrefixLen := len(bot.Config.CmdPrefix)

	// Dispatch the command iff we have our prefix (case-insensitive).
	if len(msg) > cmdPrefixLen && strings.EqualFold(bot.Config.CmdPrefix, msg[0:cmdPrefixLen]) {
		afterPrefix := msg[cmdPrefixLen:]
		cmd := strings.SplitN(afterPrefix, " ", 2)

		if f, ok := bot.cmdHandlers[strings.ToLower(cmd[0])]; ok {
			rest := ""
			if len(cmd) > 1 {
				rest = cmd[1]
			}

			if f.handler != nil {
				f.handler(s, m, rest)
			} else {
				log.Printf("Error: nil handler for command '%s'", cmd)
				s.ChannelMessageSend(m.ChannelID, "Command error - please contact support.")
			}
		}
	}
}

func (bot *Bot) onCmdHelp(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	s.ChannelMessageSend(m.ChannelID, bot.helpMessage)
}

func (bot *Bot) onCmdInvite(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	s.ChannelMessageSend(m.ChannelID, bot.inviteMessage)
}

func (bot *Bot) onCmdStartPom(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	channel, err := s.State.Channel(m.ChannelID)
	if err != nil {
		// Could not find the channel, so simply log and exit
		log.Printf("Could not find channel for ChannelID '%s'", m.ChannelID)
		return
	}

	// Make sure the user's text can't break out of our quote box.
	extra = strings.Replace(extra, "`", "", -1)

	notif := NotifyInfo{
		extra,
		m.Author.ID,
		channel.GuildID,
	}

	if bot.poms.CreateIfEmpty(m.ChannelID, pomDuration, func() { bot.onPomEnded(m.ChannelID) }, notif) {
		taskStr := "Started task"
		if len(notif.Title) > 0 {
			taskStr = fmt.Sprintf("`%s`", notif.Title)
		}
		msg := fmt.Sprintf("%s  -  **%.1f minutes** remaining!", taskStr, pomDuration.Minutes())
		s.ChannelMessageSend(m.ChannelID, msg)
	} else {
		s.ChannelMessageSend(m.ChannelID, "A Pomodoro is already running on this channel.")
	}
}

func (bot *Bot) onCmdCancelPom(s *discordgo.Session, m *discordgo.MessageCreate, extra string) {
	if notif := bot.poms.RemoveIfExists(m.ChannelID); notif != nil {
		// TODO: Use the NotifyInfo here?
		s.ChannelMessageSend(m.ChannelID, "Pomodoro cancelled!")
	} else {
		s.ChannelMessageSend(m.ChannelID, "No Pomodoro running on this channel to cancel.")
	}
}

// onPomEnded performs the notification
func (bot *Bot) onPomEnded(channelID string) {
	notif := bot.poms.RemoveIfExists(channelID)
	message := "Pomodoro ended.\n"
	messageSuffix := "Time for a short break!"

	var toMention []string

	if notif != nil {
		if len(notif.Title) > 0 {
			message = fmt.Sprintf("`%s` work cycle complete.\n", notif.Title)
		}

		user, err := bot.discord.User(notif.UserID)
		if err == nil {
			toMention = append(toMention, user.Mention())
		}
		// Doing this in a goroutine so we don't wait until the audio has been played to send the text notification.
		// This isn't required, but is my preference.
		go bot.playEndSound(*notif)

		messageBody := message + messageSuffix

		if len(toMention) > 0 {
			mentions := strings.Join(toMention, " ")
			messageBody = fmt.Sprintf("%s\n%s", messageBody, mentions)
		}

		bot.discord.ChannelMessageSend(channelID, messageBody)
	}
	// If we don't have a NotifInfo value, then the task was cancelled before we were called, so don't notify
}
<|endoftext|>


<commit_before>package colors

import "fmt"

const (
	red          string = "\x1b[91m"
	cyan         string = "\x1b[36m"
	green        string = "\x1b[32m"
	yellow       string = "\x1b[33m"
	defaultStyle string = "\x1b[0m"
	boldStyle    string = "\x1b[1m"
)

func Red(output string) string {
	return fmt.Sprintf("%s%s%s", red, output, defaultStyle)
}

func Green(output string) string {
	return fmt.Sprintf("%s%s%s", green, output, defaultStyle)
}

func Cyan(output string) string {
	return fmt.Sprintf("%s%s%s", cyan, output, defaultStyle)
}

func Yellow(output string) string {
	return fmt.Sprintf("%s%s%s", yellow, output, defaultStyle)
}

func NoColor(output string) string {
	return fmt.Sprintf("%s%s%s", defaultStyle, output, defaultStyle)
}

func Bold(output string) string {
	return fmt.Sprintf("%s%s%s", boldStyle, output, defaultStyle)
}
<commit_msg>refactor colors library<commit_after>package colors

import "fmt"

const (
	red          string = "\x1b[91m"
	cyan         string = "\x1b[36m"
	green        string = "\x1b[32m"
	yellow       string = "\x1b[33m"
	defaultStyle string = "\x1b[0m"
	boldStyle    string = "\x1b[1m"
)

func Red(output string) string {
	return colorize(output, red)
}

func Green(output string) string {
	return colorize(output, green)
}

func Cyan(output string) string {
	return colorize(output, cyan)
}

func Yellow(output string) string {
	return colorize(output, yellow)
}

func NoColor(output string) string {
	return colorize(output, defaultStyle)
}

func Bold(output string) string {
	return colorize(output, boldStyle)
}

func colorize(output string, color string) string {
	return fmt.Sprintf("%s%s%s", color, output, defaultStyle)
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"os"
	"strings"
	"github.com/jingweno/gh/utils"
)

var cmdPush = &Command{
	Run: push,
	GitExtension: true,
	Usage: "push REMOTE-1,REMOTE-2,...,REMOTE-N [REF]",
	Short: "Update remote refs along with associated objects",
	Long: `Push REF to each of REMOTE-1 through REMOTE-N by executing  mul-
tiple git push commands.`,
}

/**
 $ git push origin,staging,qa bert_timeout
 > git push origin bert_timeout
 > git push staging bert_timeout
 > git push qa bert_timeout
**/

func push (command *Command, args *Args) {
	if !args.IsParamsEmpty() {
		pushToEveryRemote(args)
	}
}

func pushToEveryRemote (args *Args) {
	remotes, idx := getRemotes(args)
	for _, i := range remotes {
		copyArgs := args
		copyArgs.ReplaceParam(idx, i)
		if !args.Noop {
			err := copyArgs.ToCmd().Exec()
			utils.Check(err)
		} else {
			fmt.Printf("it would run `git push %s`\n", strings.Join(copyArgs.Params, " "))
		}
	}

	fixHelp(args)
}

func getRemotes(args *Args) (remotes []string, idx int) {
	for a, i := range args.Params {
		if !strings.HasPrefix(i, "-") {
			remotes = strings.Split(i, ",")
			idx = a
			return
		}
	}

	return
}

func fixHelp(args *Args) {
	if i := args.IndexOfParam("--help"); i != -1 && args.ParamsSize() == 1 {
		args.Params = []string{"--help"}
		return
	}

	os.Exit(0)
}
<commit_msg>Changed variable names<commit_after>package commands

import (
	"fmt"
	"os"
	"strings"
	"github.com/jingweno/gh/utils"
)

var cmdPush = &Command{
	Run: push,
	GitExtension: true,
	Usage: "push REMOTE-1,REMOTE-2,...,REMOTE-N [REF]",
	Short: "Update remote refs along with associated objects",
	Long: `Push REF to each of REMOTE-1 through REMOTE-N by executing  mul-
tiple git push commands.`,
}

/**
 $ git push origin,staging,qa bert_timeout
 > git push origin bert_timeout
 > git push staging bert_timeout
 > git push qa bert_timeout
**/

func push (command *Command, args *Args) {
	if !args.IsParamsEmpty() {
		pushToEveryRemote(args)
	}
}

func pushToEveryRemote (args *Args) {
	remotes, idx := getRemotes(args)
	for _, a := range remotes {
		copyArgs := args
		copyArgs.ReplaceParam(idx, a)
		if !args.Noop {
			err := copyArgs.ToCmd().Exec()
			utils.Check(err)
		} else {
			fmt.Printf("it would run `git push %s`\n", strings.Join(copyArgs.Params, " "))
		}
	}

	fixHelp(args)
}

func getRemotes(args *Args) (remotes []string, idx int) {
	for i, a := range args.Params {
		if !strings.HasPrefix(a, "-") {
			remotes = strings.Split(a, ",")
			idx = i
			return
		}
	}

	return
}

func fixHelp(args *Args) {
	if i := args.IndexOfParam("--help"); i != -1 && args.ParamsSize() == 1 {
		args.Params = []string{"--help"}
		return
	}

	os.Exit(0)
}
<|endoftext|>


<commit_before>package main

import (
	"natyla"
    "testing"
)

/*
Test the function that convert string numbers to integers
*/
func TestAtoi(t *testing.T) {
	result := natyla.Atoi("5")
	if result !=5 {
		t.Errorf("Atoi(%s) returned %d, expected %d", "4", result, 5)
	}
}
<commit_msg>CI travis config<commit_after>package main

import (
	"github.com/Fersca/natyla/src/natyla"
    "testing"
)

/*
Test the function that convert string numbers to integers
*/
func TestAtoi(t *testing.T) {
	result := natyla.Atoi("5")
	if result !=5 {
		t.Errorf("Atoi(%s) returned %d, expected %d", "4", result, 5)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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

    https://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 scheduledsparkapplication

import (
	"fmt"
	"reflect"
	"time"

	"github.com/golang/glog"
	"github.com/robfig/cron"

	apiextensionsclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/clock"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/retry"
	"k8s.io/client-go/util/workqueue"

	"k8s.io/spark-on-k8s-operator/pkg/apis/sparkoperator.k8s.io/v1alpha1"
	crdclientset "k8s.io/spark-on-k8s-operator/pkg/client/clientset/versioned"
	crdscheme "k8s.io/spark-on-k8s-operator/pkg/client/clientset/versioned/scheme"
	crdinformers "k8s.io/spark-on-k8s-operator/pkg/client/informers/externalversions"
	crdlisters "k8s.io/spark-on-k8s-operator/pkg/client/listers/sparkoperator.k8s.io/v1alpha1"
)

var (
	keyFunc = cache.DeletionHandlingMetaNamespaceKeyFunc
)

type Controller struct {
	crdClient        crdclientset.Interface
	kubeClient       kubernetes.Interface
	extensionsClient apiextensionsclient.Interface
	queue            workqueue.RateLimitingInterface
	cacheSynced      cache.InformerSynced
	lister           crdlisters.ScheduledSparkApplicationLister
	clock            clock.Clock
}

func NewController(
	crdClient crdclientset.Interface,
	kubeClient kubernetes.Interface,
	extensionsClient apiextensionsclient.Interface,
	informerFactory crdinformers.SharedInformerFactory,
	clock clock.Clock) *Controller {
	crdscheme.AddToScheme(scheme.Scheme)

	queue := workqueue.NewNamedRateLimitingQueue(workqueue.DefaultControllerRateLimiter(),
		"scheduled-spark-application-controller")

	controller := &Controller{
		crdClient:        crdClient,
		kubeClient:       kubeClient,
		extensionsClient: extensionsClient,
		queue:            queue,
		clock:            clock,
	}

	informer := informerFactory.Sparkoperator().V1alpha1().ScheduledSparkApplications()
	informer.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc:    controller.onAdd,
		UpdateFunc: controller.onUpdate,
		DeleteFunc: controller.onDelete,
	})
	controller.cacheSynced = informer.Informer().HasSynced
	controller.lister = informer.Lister()

	return controller
}

func (c *Controller) Start(workers int, stopCh <-chan struct{}) error {
	glog.Info("Starting the ScheduledSparkApplication controller")

	if !cache.WaitForCacheSync(stopCh, c.cacheSynced) {
		return fmt.Errorf("timed out waiting for cache to sync")
	}

	glog.Info("Starting the workers of the ScheduledSparkApplication controller")
	for i := 0; i < workers; i++ {
		// runWorker will loop until "something bad" happens. Until will then rekick
		// the worker after one second.
		go wait.Until(c.runWorker, time.Second, stopCh)
	}

	return nil
}

func (c *Controller) Stop() {
	glog.Info("Stopping the ScheduledSparkApplication controller")
	c.queue.ShutDown()
}

func (c *Controller) runWorker() {
	defer utilruntime.HandleCrash()
	for c.processNextItem() {
	}
}

func (c *Controller) processNextItem() bool {
	key, quit := c.queue.Get()
	if quit {
		return false
	}
	defer c.queue.Done(key)

	err := c.syncScheduledSparkApplication(key.(string))
	if err == nil {
		// Successfully processed the key or the key was not found so tell the queue to stop tracking
		// history for your key. This will reset things like failure counts for per-item rate limiting.
		c.queue.Forget(key)
		return true
	}

	// There was a failure so be sure to report it. This method allows for pluggable error handling
	// which can be used for things like cluster-monitoring
	utilruntime.HandleError(fmt.Errorf("failed to sync ScheduledSparkApplication %q: %v", key, err))
	// Since we failed, we should requeue the item to work on later.  This method will add a backoff
	// to avoid hot-looping on particular items (they're probably still not going to work right away)
	// and overall controller protection (everything I've done is broken, this controller needs to
	// calm down or it can starve other useful work) cases.
	c.queue.AddRateLimited(key)

	return true
}

func (c *Controller) syncScheduledSparkApplication(key string) error {
	namespace, name, err := cache.SplitMetaNamespaceKey(key)
	if err != nil {
		return err
	}
	app, err := c.lister.ScheduledSparkApplications(namespace).Get(name)
	if err != nil {
		return err
	}

	if app.Spec.Suspend != nil && *app.Spec.Suspend {
		return nil
	}

	glog.V(2).Infof("Syncing ScheduledSparkApplication %s", app.Name)
	status := app.Status.DeepCopy()
	schedule, err := cron.ParseStandard(app.Spec.Schedule)
	if err != nil {
		glog.Errorf("failed to parse schedule %s of %s: %v", app.Spec.Schedule, app.Name, err)
		status.ScheduleState = v1alpha1.FailedValidationState
		status.Reason = err.Error()
	} else {
		status.ScheduleState = v1alpha1.ScheduledState
		nextRunTime := status.NextRun.Time
		if nextRunTime.IsZero() {
			nextRunTime = schedule.Next(status.LastRun.Time)
			status.NextRun = metav1.NewTime(nextRunTime)
		}
		now := c.clock.Now()
		if nextRunTime.Before(now) {
			// The next run is due. Check if this is the first run of the application.
			if len(status.PastRunNames) == 0 {
				// This is the first run of the application.
				if err = c.startNextRun(app, status, schedule); err != nil {
					return err
				}
			} else {
				// Check if the condition for starting the next run is satisfied.
				ok, err := c.shouldStartNextRun(app)
				if err != nil {
					return err
				}
				if ok {
					if err = c.startNextRun(app, status, schedule); err != nil {
						return err
					}
				}
			}
		}
	}

	return c.updateScheduledSparkApplicationStatus(app, status)
}

func (c *Controller) onAdd(obj interface{}) {
	c.enqueue(obj)
}

func (c *Controller) onUpdate(oldObj, newObj interface{}) {
	c.enqueue(newObj)
}

func (c *Controller) onDelete(obj interface{}) {
	c.dequeue(obj)
}

func (c *Controller) enqueue(obj interface{}) {
	key, err := keyFunc(obj)
	if err != nil {
		glog.Errorf("failed to get key for %v: %v", obj, err)
		return
	}

	c.queue.AddRateLimited(key)
}

func (c *Controller) dequeue(obj interface{}) {
	key, err := keyFunc(obj)
	if err != nil {
		glog.Errorf("failed to get key for %v: %v", obj, err)
		return
	}

	c.queue.Forget(key)
	c.queue.Done(key)
}

func (c *Controller) createSparkApplication(
	scheduledApp *v1alpha1.ScheduledSparkApplication, t time.Time) (string, error) {
	app := &v1alpha1.SparkApplication{}
	app.Spec = scheduledApp.Spec.Template
	app.Name = fmt.Sprintf("%s-%d", scheduledApp.Name, t.UnixNano())
	app.OwnerReferences = append(app.OwnerReferences, metav1.OwnerReference{
		APIVersion: v1alpha1.SchemeGroupVersion.String(),
		Kind:       reflect.TypeOf(v1alpha1.ScheduledSparkApplication{}).Name(),
		Name:       scheduledApp.Name,
		UID:        scheduledApp.UID,
	})
	_, err := c.crdClient.SparkoperatorV1alpha1().SparkApplications(scheduledApp.Namespace).Create(app)
	if err != nil {
		return "", err
	}
	return app.Name, nil
}

func (c *Controller) shouldStartNextRun(app *v1alpha1.ScheduledSparkApplication) (bool, error) {
	switch app.Spec.ConcurrencyPolicy {
	case v1alpha1.ConcurrencyAllow:
		return true, nil
	case v1alpha1.ConcurrencyForbid:
		finished, _, err := c.hasLastRunFinished(app.Namespace, app.Status.PastRunNames[0])
		if err != nil {
			return false, err
		}
		return finished, nil
	case v1alpha1.ConcurrencyReplace:
		if err := c.killLastRunIfNotFinished(app.Namespace, app.Status.PastRunNames[0]); err != nil {
			return false, err
		}
		return true, nil
	}
	return true, nil
}

func (c *Controller) startNextRun(
	app *v1alpha1.ScheduledSparkApplication,
	status *v1alpha1.ScheduledSparkApplicationStatus,
	schedule cron.Schedule) error {
	glog.Infof("Next run of %s is due, creating a new SparkApplication instance", app.Name)
	status.LastRun = metav1.Now()
	name, err := c.createSparkApplication(app, status.LastRun.Time)
	if err != nil {
		glog.Errorf("failed to create a SparkApplication instance for %s: %v", app.Name, err)
		return err
	}
	status.NextRun = metav1.NewTime(schedule.Next(status.LastRun.Time))

	var limit int32 = 1
	if app.Spec.RunHistoryLimit != nil {
		limit = *app.Spec.RunHistoryLimit
	}

	rest := status.PastRunNames
	var toDelete []string
	if int32(len(status.PastRunNames)) >= limit {
		rest = status.PastRunNames[:limit-1]
		toDelete = status.PastRunNames[limit-1:]
	}
	// Pre-append the name of the latest run.
	status.PastRunNames = append([]string{name}, rest...)

	namespace := app.Namespace
	// Delete runs that should no longer be kept.
	for _, name := range toDelete {
		c.crdClient.SparkoperatorV1alpha1().SparkApplications(namespace).Delete(name, metav1.NewDeleteOptions(0))
	}

	return nil
}

func (c *Controller) hasLastRunFinished(
	namespace string,
	lastRunName string) (bool, *v1alpha1.SparkApplication, error) {
	app, err := c.crdClient.SparkoperatorV1alpha1().SparkApplications(namespace).Get(lastRunName, metav1.GetOptions{})
	if err != nil {
		if errors.IsNotFound(err) {
			return true, nil, nil
		}
		return false, nil, err
	}

	return app.Status.AppState.State == v1alpha1.CompletedState ||
		app.Status.AppState.State == v1alpha1.FailedState, app, nil
}

func (c *Controller) killLastRunIfNotFinished(namespace string, lastRunName string) error {
	finished, app, err := c.hasLastRunFinished(namespace, lastRunName)
	if err != nil {
		return err
	}
	if app == nil || finished {
		return nil
	}

	// Delete the driver pod of the last run if applicable.
	if app.Status.DriverInfo.PodName != "" {
		if err = c.kubeClient.CoreV1().Pods(namespace).Delete(app.Status.DriverInfo.PodName,
			metav1.NewDeleteOptions(0)); err != nil {
			return err
		}
	}

	// Delete the SparkApplication object of the last run.
	if err = c.crdClient.SparkoperatorV1alpha1().SparkApplications(namespace).Delete(lastRunName,
		metav1.NewDeleteOptions(0)); err != nil {
		return err
	}

	return nil
}

func (c *Controller) updateScheduledSparkApplicationStatus(
	app *v1alpha1.ScheduledSparkApplication,
	newStatus *v1alpha1.ScheduledSparkApplicationStatus) error {
	// If the status has not changed, do not perform an update.
	if isStatusEqual(newStatus, &app.Status) {
		return nil
	}

	toUpdate := app.DeepCopy()
	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
		toUpdate.Status = *newStatus
		_, updateErr := c.crdClient.SparkoperatorV1alpha1().ScheduledSparkApplications(toUpdate.Namespace).Update(
			toUpdate)
		if updateErr == nil {
			return nil
		}

		result, err := c.crdClient.SparkoperatorV1alpha1().ScheduledSparkApplications(toUpdate.Namespace).Get(
			toUpdate.Name, metav1.GetOptions{})
		if err != nil {
			return err
		}
		toUpdate = result

		return updateErr
	})
}

func isStatusEqual(newStatus, currentStatus *v1alpha1.ScheduledSparkApplicationStatus) bool {
	return newStatus.ScheduleState == currentStatus.ScheduleState &&
		newStatus.LastRun == currentStatus.LastRun &&
		newStatus.NextRun == currentStatus.NextRun &&
		reflect.DeepEqual(newStatus.PastRunNames, currentStatus.PastRunNames) &&
		newStatus.Reason == currentStatus.Reason
}
<commit_msg>Set LastRun only if the controller successfully creates the SparkApplication<commit_after>/*
Copyright 2018 Google LLC

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

    https://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 scheduledsparkapplication

import (
	"fmt"
	"reflect"
	"time"

	"github.com/golang/glog"
	"github.com/robfig/cron"

	apiextensionsclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/clock"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/retry"
	"k8s.io/client-go/util/workqueue"

	"k8s.io/spark-on-k8s-operator/pkg/apis/sparkoperator.k8s.io/v1alpha1"
	crdclientset "k8s.io/spark-on-k8s-operator/pkg/client/clientset/versioned"
	crdscheme "k8s.io/spark-on-k8s-operator/pkg/client/clientset/versioned/scheme"
	crdinformers "k8s.io/spark-on-k8s-operator/pkg/client/informers/externalversions"
	crdlisters "k8s.io/spark-on-k8s-operator/pkg/client/listers/sparkoperator.k8s.io/v1alpha1"
)

var (
	keyFunc = cache.DeletionHandlingMetaNamespaceKeyFunc
)

type Controller struct {
	crdClient        crdclientset.Interface
	kubeClient       kubernetes.Interface
	extensionsClient apiextensionsclient.Interface
	queue            workqueue.RateLimitingInterface
	cacheSynced      cache.InformerSynced
	lister           crdlisters.ScheduledSparkApplicationLister
	clock            clock.Clock
}

func NewController(
	crdClient crdclientset.Interface,
	kubeClient kubernetes.Interface,
	extensionsClient apiextensionsclient.Interface,
	informerFactory crdinformers.SharedInformerFactory,
	clock clock.Clock) *Controller {
	crdscheme.AddToScheme(scheme.Scheme)

	queue := workqueue.NewNamedRateLimitingQueue(workqueue.DefaultControllerRateLimiter(),
		"scheduled-spark-application-controller")

	controller := &Controller{
		crdClient:        crdClient,
		kubeClient:       kubeClient,
		extensionsClient: extensionsClient,
		queue:            queue,
		clock:            clock,
	}

	informer := informerFactory.Sparkoperator().V1alpha1().ScheduledSparkApplications()
	informer.Informer().AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc:    controller.onAdd,
		UpdateFunc: controller.onUpdate,
		DeleteFunc: controller.onDelete,
	})
	controller.cacheSynced = informer.Informer().HasSynced
	controller.lister = informer.Lister()

	return controller
}

func (c *Controller) Start(workers int, stopCh <-chan struct{}) error {
	glog.Info("Starting the ScheduledSparkApplication controller")

	if !cache.WaitForCacheSync(stopCh, c.cacheSynced) {
		return fmt.Errorf("timed out waiting for cache to sync")
	}

	glog.Info("Starting the workers of the ScheduledSparkApplication controller")
	for i := 0; i < workers; i++ {
		// runWorker will loop until "something bad" happens. Until will then rekick
		// the worker after one second.
		go wait.Until(c.runWorker, time.Second, stopCh)
	}

	return nil
}

func (c *Controller) Stop() {
	glog.Info("Stopping the ScheduledSparkApplication controller")
	c.queue.ShutDown()
}

func (c *Controller) runWorker() {
	defer utilruntime.HandleCrash()
	for c.processNextItem() {
	}
}

func (c *Controller) processNextItem() bool {
	key, quit := c.queue.Get()
	if quit {
		return false
	}
	defer c.queue.Done(key)

	err := c.syncScheduledSparkApplication(key.(string))
	if err == nil {
		// Successfully processed the key or the key was not found so tell the queue to stop tracking
		// history for your key. This will reset things like failure counts for per-item rate limiting.
		c.queue.Forget(key)
		return true
	}

	// There was a failure so be sure to report it. This method allows for pluggable error handling
	// which can be used for things like cluster-monitoring
	utilruntime.HandleError(fmt.Errorf("failed to sync ScheduledSparkApplication %q: %v", key, err))
	// Since we failed, we should requeue the item to work on later.  This method will add a backoff
	// to avoid hot-looping on particular items (they're probably still not going to work right away)
	// and overall controller protection (everything I've done is broken, this controller needs to
	// calm down or it can starve other useful work) cases.
	c.queue.AddRateLimited(key)

	return true
}

func (c *Controller) syncScheduledSparkApplication(key string) error {
	namespace, name, err := cache.SplitMetaNamespaceKey(key)
	if err != nil {
		return err
	}
	app, err := c.lister.ScheduledSparkApplications(namespace).Get(name)
	if err != nil {
		return err
	}

	if app.Spec.Suspend != nil && *app.Spec.Suspend {
		return nil
	}

	glog.V(2).Infof("Syncing ScheduledSparkApplication %s", app.Name)
	status := app.Status.DeepCopy()
	schedule, err := cron.ParseStandard(app.Spec.Schedule)
	if err != nil {
		glog.Errorf("failed to parse schedule %s of %s: %v", app.Spec.Schedule, app.Name, err)
		status.ScheduleState = v1alpha1.FailedValidationState
		status.Reason = err.Error()
	} else {
		status.ScheduleState = v1alpha1.ScheduledState
		nextRunTime := status.NextRun.Time
		if nextRunTime.IsZero() {
			nextRunTime = schedule.Next(status.LastRun.Time)
			status.NextRun = metav1.NewTime(nextRunTime)
		}
		now := c.clock.Now()
		if nextRunTime.Before(now) {
			// The next run is due. Check if this is the first run of the application.
			if len(status.PastRunNames) == 0 {
				// This is the first run of the application.
				if err = c.startNextRun(app, status, schedule); err != nil {
					return err
				}
			} else {
				// Check if the condition for starting the next run is satisfied.
				ok, err := c.shouldStartNextRun(app)
				if err != nil {
					return err
				}
				if ok {
					if err = c.startNextRun(app, status, schedule); err != nil {
						return err
					}
				}
			}
		}
	}

	return c.updateScheduledSparkApplicationStatus(app, status)
}

func (c *Controller) onAdd(obj interface{}) {
	c.enqueue(obj)
}

func (c *Controller) onUpdate(oldObj, newObj interface{}) {
	c.enqueue(newObj)
}

func (c *Controller) onDelete(obj interface{}) {
	c.dequeue(obj)
}

func (c *Controller) enqueue(obj interface{}) {
	key, err := keyFunc(obj)
	if err != nil {
		glog.Errorf("failed to get key for %v: %v", obj, err)
		return
	}

	c.queue.AddRateLimited(key)
}

func (c *Controller) dequeue(obj interface{}) {
	key, err := keyFunc(obj)
	if err != nil {
		glog.Errorf("failed to get key for %v: %v", obj, err)
		return
	}

	c.queue.Forget(key)
	c.queue.Done(key)
}

func (c *Controller) createSparkApplication(
	scheduledApp *v1alpha1.ScheduledSparkApplication, t time.Time) (string, error) {
	app := &v1alpha1.SparkApplication{}
	app.Spec = scheduledApp.Spec.Template
	app.Name = fmt.Sprintf("%s-%d", scheduledApp.Name, t.UnixNano())
	app.OwnerReferences = append(app.OwnerReferences, metav1.OwnerReference{
		APIVersion: v1alpha1.SchemeGroupVersion.String(),
		Kind:       reflect.TypeOf(v1alpha1.ScheduledSparkApplication{}).Name(),
		Name:       scheduledApp.Name,
		UID:        scheduledApp.UID,
	})
	_, err := c.crdClient.SparkoperatorV1alpha1().SparkApplications(scheduledApp.Namespace).Create(app)
	if err != nil {
		return "", err
	}
	return app.Name, nil
}

func (c *Controller) shouldStartNextRun(app *v1alpha1.ScheduledSparkApplication) (bool, error) {
	switch app.Spec.ConcurrencyPolicy {
	case v1alpha1.ConcurrencyAllow:
		return true, nil
	case v1alpha1.ConcurrencyForbid:
		finished, _, err := c.hasLastRunFinished(app.Namespace, app.Status.PastRunNames[0])
		if err != nil {
			return false, err
		}
		return finished, nil
	case v1alpha1.ConcurrencyReplace:
		if err := c.killLastRunIfNotFinished(app.Namespace, app.Status.PastRunNames[0]); err != nil {
			return false, err
		}
		return true, nil
	}
	return true, nil
}

func (c *Controller) startNextRun(
	app *v1alpha1.ScheduledSparkApplication,
	status *v1alpha1.ScheduledSparkApplicationStatus,
	schedule cron.Schedule) error {
	glog.Infof("Next run of %s is due, creating a new SparkApplication instance", app.Name)
	now := metav1.Now()
	name, err := c.createSparkApplication(app, now.Time)
	if err != nil {
		glog.Errorf("failed to create a SparkApplication instance for %s: %v", app.Name, err)
		return err
	}

	status.LastRun = now
	status.NextRun = metav1.NewTime(schedule.Next(status.LastRun.Time))

	var limit int32 = 1
	if app.Spec.RunHistoryLimit != nil {
		limit = *app.Spec.RunHistoryLimit
	}

	rest := status.PastRunNames
	var toDelete []string
	if int32(len(status.PastRunNames)) >= limit {
		rest = status.PastRunNames[:limit-1]
		toDelete = status.PastRunNames[limit-1:]
	}
	// Pre-append the name of the latest run.
	status.PastRunNames = append([]string{name}, rest...)

	namespace := app.Namespace
	// Delete runs that should no longer be kept.
	for _, name := range toDelete {
		c.crdClient.SparkoperatorV1alpha1().SparkApplications(namespace).Delete(name, metav1.NewDeleteOptions(0))
	}

	return nil
}

func (c *Controller) hasLastRunFinished(
	namespace string,
	lastRunName string) (bool, *v1alpha1.SparkApplication, error) {
	app, err := c.crdClient.SparkoperatorV1alpha1().SparkApplications(namespace).Get(lastRunName, metav1.GetOptions{})
	if err != nil {
		if errors.IsNotFound(err) {
			return true, nil, nil
		}
		return false, nil, err
	}

	return app.Status.AppState.State == v1alpha1.CompletedState ||
		app.Status.AppState.State == v1alpha1.FailedState, app, nil
}

func (c *Controller) killLastRunIfNotFinished(namespace string, lastRunName string) error {
	finished, app, err := c.hasLastRunFinished(namespace, lastRunName)
	if err != nil {
		return err
	}
	if app == nil || finished {
		return nil
	}

	// Delete the driver pod of the last run if applicable.
	if app.Status.DriverInfo.PodName != "" {
		if err = c.kubeClient.CoreV1().Pods(namespace).Delete(app.Status.DriverInfo.PodName,
			metav1.NewDeleteOptions(0)); err != nil {
			return err
		}
	}

	// Delete the SparkApplication object of the last run.
	if err = c.crdClient.SparkoperatorV1alpha1().SparkApplications(namespace).Delete(lastRunName,
		metav1.NewDeleteOptions(0)); err != nil {
		return err
	}

	return nil
}

func (c *Controller) updateScheduledSparkApplicationStatus(
	app *v1alpha1.ScheduledSparkApplication,
	newStatus *v1alpha1.ScheduledSparkApplicationStatus) error {
	// If the status has not changed, do not perform an update.
	if isStatusEqual(newStatus, &app.Status) {
		return nil
	}

	toUpdate := app.DeepCopy()
	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
		toUpdate.Status = *newStatus
		_, updateErr := c.crdClient.SparkoperatorV1alpha1().ScheduledSparkApplications(toUpdate.Namespace).Update(
			toUpdate)
		if updateErr == nil {
			return nil
		}

		result, err := c.crdClient.SparkoperatorV1alpha1().ScheduledSparkApplications(toUpdate.Namespace).Get(
			toUpdate.Name, metav1.GetOptions{})
		if err != nil {
			return err
		}
		toUpdate = result

		return updateErr
	})
}

func isStatusEqual(newStatus, currentStatus *v1alpha1.ScheduledSparkApplicationStatus) bool {
	return newStatus.ScheduleState == currentStatus.ScheduleState &&
		newStatus.LastRun == currentStatus.LastRun &&
		newStatus.NextRun == currentStatus.NextRun &&
		reflect.DeepEqual(newStatus.PastRunNames, currentStatus.PastRunNames) &&
		newStatus.Reason == currentStatus.Reason
}
<|endoftext|>


<commit_before>package moldova

import (
	"bytes"
	"errors"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"
)

type TestComparator func(string) error

type TestCase struct {
	Template     string
	Comparator   TestComparator
	ParseFailure bool
	WriteFailure bool
}

var GUIDCases = []TestCase{
	{
		Template: "{guid}",
		Comparator: func(s string) error {
			p := strings.Split(s, "-")
			if len(p) == 5 &&
				len(p[0]) == 8 &&
				len(p[1]) == len(p[2]) && len(p[2]) == len(p[3]) && len(p[3]) == 4 &&
				len(p[4]) == 12 {
				return nil
			}
			return errors.New("Guid not in correct format: " + s)
		},
	},
	{
		Template: "{guid}@{guid:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Guid at position 1 not equal to guid at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{guid}@{guid:ordinal:1}",
		WriteFailure: true,
	},
}

var NowCases = []TestCase{
	{
		// There is no proper deterministic way to test what the value of now is, without
		// something like rubys timecop (but the go-equivalent is not viable) or relying
		// on luck, which will run out if tests are run at just the wrong moment.
		// Therefore, for the basic test, i'm just asserting nothing went wrong for now.
		Template: "{now}",
		Comparator: func(s string) error {
			if len(s) > 0 {
				return nil
			}
			return errors.New("Now not in correct format: " + s)
		},
	},
	{
		Template: "{now}@{now:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Now at position 1 not equal to now at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{now}@{now:ordinal:1}",
		WriteFailure: true,
	},
}

var TimeCases = []TestCase{
	{
		Template: "{time:min:1|max:1|format:simple|zone:EST}",
		Comparator: func(s string) error {
			if s == "1969-12-31 19:00:01" {
				return nil
			}
			return errors.New("Time value was not the expected value")
		},
	},
	{
		Template: "{time:min:1|max:1|format:simpletz|zone:EST}",
		Comparator: func(s string) error {
			if s == "1969-12-31 19:00:01 -0500" {
				return nil
			}
			return errors.New("Time value was not the expected value")
		},
	},
	{
		Template: "{time:min:1|max:1|format:2006//01//02@@15_04_05}|zone:EST",
		Comparator: func(s string) error {
			if s == "1970//01//01@@00_00_01" {
				return nil
			}
			return errors.New("Time value was not the expected value")
		},
	},
	{
		Template: "{time}@{time:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Time at position 1 not equal to time at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{time}@{time:ordinal:1}",
		WriteFailure: true,
	},
}

var CountryCases = []TestCase{
	{
		Template: "{country}",
		Comparator: func(s string) error {
			// TODO better check here in case we ever support different types of country codes
			if len(s) == 2 {
				return nil
			}
			return errors.New("Invalid country code generated somehow")
		},
	},
	{
		Template: "{country:case:up}",
		Comparator: func(s string) error {
			// Since I can't know which country comes out, i'll invert the result
			// If the ToLowered result is not the same as the original result, we know
			// that the original was successfully output in upper case
			if strings.ToLower(s) != s {
				return nil
			}
			return errors.New("Country was returned in lowercase, but was requested in uppercase")
		},
	},
	{
		Template: "{country:case:down}",
		Comparator: func(s string) error {
			// Since I can't know which country comes out, i'll invert the result
			// If the ToLowered result is not the same as the original result, we know
			// that the original was successfully output in upper case
			if strings.ToUpper(s) != s {
				return nil
			}
			return errors.New("Country was returned in uppercase, but was requested in lowercase")
		},
	},
	{
		Template: "{country}@{country:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Country at position 1 not equal to country at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{country}@{country:ordinal:1}",
		WriteFailure: true,
	},
}

// Placeholders
var FloatCases = []TestCase{
	{
		Template: "{float}",
		Comparator: func(s string) error {
			i, err := strconv.ParseFloat(s, 64)
			if err != nil {
				return err
			}
			if i >= 0.0 && i <= 100.0 {
				return nil
			}
			return errors.New("Float out of range for default min/max values")
		},
	},
	{
		Template: "{float:max:5000.0|min:4999.0}",
		Comparator: func(s string) error {
			i, err := strconv.ParseFloat(s, 64)
			if err != nil {
				return err
			}
			if i >= 4999.0 && i <= 5000.0 {
				return nil
			}
			return errors.New("Float out of range for custom min/max values")
		},
	},
	{
		Template: "{float:max:-5000.0|min:-5001.0}",
		Comparator: func(s string) error {
			i, err := strconv.ParseFloat(s, 64)
			if err != nil {
				return err
			}
			if i >= -5001.0 && i <= -5000.0 {
				return nil
			}
			return errors.New("Float out of range for custom min/max values")
		},
	},
	{
		Template: "{float}@{float:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Float at position 1 not equal to Float at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{float}@{float:ordinal:1}",
		WriteFailure: true,
	},
}

var IntegerCases = []TestCase{
	{
		Template: "{int}",
		Comparator: func(s string) error {
			i, err := strconv.Atoi(s)
			if err != nil {
				return err
			}
			if i >= 0 && i <= 100 {
				return nil
			}
			return errors.New("Int out of range for default min/max values")
		},
	},
	{
		Template: "{int:max:5000|min:4999}",
		Comparator: func(s string) error {
			i, err := strconv.Atoi(s)
			if err != nil {
				return err
			}
			if i >= 4999 && i <= 5000 {
				return nil
			}
			return errors.New("Int out of range for custom min/max values")
		},
	},
	{
		Template: "{int:max:-5000|min:-5001}",
		Comparator: func(s string) error {
			i, err := strconv.Atoi(s)
			if err != nil {
				return err
			}
			if i >= -5001 && i <= -5000 {
				return nil
			}
			return errors.New("Int out of range for custom min/max values")
		},
	},
	{
		Template: "{int}@{int:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Int at position 1 not equal to int at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{int}@{int:ordinal:1}",
		WriteFailure: true,
	},
}

var AllCases = [][]TestCase{
	GUIDCases,
	NowCases,
	// The TimeCases collection does not run correctly on Travis-ci due to poor
	// assumptions baked into the tests, revolving around what time zone the machine
	// running the test is on. I've got a fix for this in the roadmap
	TimeCases,
	CountryCases,
	FloatCases,
	IntegerCases,
}

// TODO Test each random function individually, under a number of inputs to make supported
// all the options behave as expected.

func TestMain(m *testing.M) {
	rand.Seed(time.Now().Unix())
	os.Exit(m.Run())
}

func TestAllCases(t *testing.T) {
	// TODO The library should be threadsafe, I should go wide here to run all specs
	// in parallel, like the natural tests would be. Channel + waitgroup to collect
	// and report on errors once they all finish
	for _, cs := range AllCases {
		for _, c := range cs {
			cs, err := BuildCallstack(c.Template)
			// If we get an error and weren't expecting it
			// Or, if we didn't get one but were expecting it
			if err != nil && !c.ParseFailure {
				t.Error(err)
			} else if err == nil && c.ParseFailure {
				t.Error("Expected to encounter Parse Failure, but did not for Test Case ", c.Template)
			}

			result := &bytes.Buffer{}
			err = cs.Write(result)

			// If we get an error and weren't expecting it
			// Or, if we didn't get one but were expecting it
			if err != nil && !c.WriteFailure {
				t.Error(err)
			} else if err == nil && c.ParseFailure {
				t.Error("Expected to encounter Write Failure, but did not for Test Case ", c.Template)
			}

			if c.Comparator != nil {
				if err := c.Comparator(result.String()); err != nil {
					t.Error(err)
				}
			}
		}
	}
}

func BenchmarkWrites(b *testing.B) {
	template := "INSERT INTO floof VALUES ('{guid}','{time},'{guid:ordinal:0}','{country}',{int:min:-2000|max:0},{int:min:100|max:1000},{float:min:-1000.0|max:-540.0},{int:min:1|max:40},'{now}','{now:ordinal:0}','{unicode:length:2|case:up}',NULL,-3)"
	var cs *Callstack
	var err error
	if cs, err = BuildCallstack(template); err != nil {
		b.Error(err)
	}
	for n := 0; n < b.N; n++ {
		result := &bytes.Buffer{}
		err = cs.Write(result)
		if err != nil {
			b.Error(err)
		}
	}
}
<commit_msg>Removing old comment<commit_after>package moldova

import (
	"bytes"
	"errors"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"
)

type TestComparator func(string) error

type TestCase struct {
	Template     string
	Comparator   TestComparator
	ParseFailure bool
	WriteFailure bool
}

var GUIDCases = []TestCase{
	{
		Template: "{guid}",
		Comparator: func(s string) error {
			p := strings.Split(s, "-")
			if len(p) == 5 &&
				len(p[0]) == 8 &&
				len(p[1]) == len(p[2]) && len(p[2]) == len(p[3]) && len(p[3]) == 4 &&
				len(p[4]) == 12 {
				return nil
			}
			return errors.New("Guid not in correct format: " + s)
		},
	},
	{
		Template: "{guid}@{guid:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Guid at position 1 not equal to guid at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{guid}@{guid:ordinal:1}",
		WriteFailure: true,
	},
}

var NowCases = []TestCase{
	{
		// There is no proper deterministic way to test what the value of now is, without
		// something like rubys timecop (but the go-equivalent is not viable) or relying
		// on luck, which will run out if tests are run at just the wrong moment.
		// Therefore, for the basic test, i'm just asserting nothing went wrong for now.
		Template: "{now}",
		Comparator: func(s string) error {
			if len(s) > 0 {
				return nil
			}
			return errors.New("Now not in correct format: " + s)
		},
	},
	{
		Template: "{now}@{now:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Now at position 1 not equal to now at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{now}@{now:ordinal:1}",
		WriteFailure: true,
	},
}

var TimeCases = []TestCase{
	{
		Template: "{time:min:1|max:1|format:simple|zone:EST}",
		Comparator: func(s string) error {
			if s == "1969-12-31 19:00:01" {
				return nil
			}
			return errors.New("Time value was not the expected value")
		},
	},
	{
		Template: "{time:min:1|max:1|format:simpletz|zone:EST}",
		Comparator: func(s string) error {
			if s == "1969-12-31 19:00:01 -0500" {
				return nil
			}
			return errors.New("Time value was not the expected value")
		},
	},
	{
		Template: "{time:min:1|max:1|format:2006//01//02@@15_04_05}|zone:EST",
		Comparator: func(s string) error {
			if s == "1970//01//01@@00_00_01" {
				return nil
			}
			return errors.New("Time value was not the expected value")
		},
	},
	{
		Template: "{time}@{time:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Time at position 1 not equal to time at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{time}@{time:ordinal:1}",
		WriteFailure: true,
	},
}

var CountryCases = []TestCase{
	{
		Template: "{country}",
		Comparator: func(s string) error {
			// TODO better check here in case we ever support different types of country codes
			if len(s) == 2 {
				return nil
			}
			return errors.New("Invalid country code generated somehow")
		},
	},
	{
		Template: "{country:case:up}",
		Comparator: func(s string) error {
			// Since I can't know which country comes out, i'll invert the result
			// If the ToLowered result is not the same as the original result, we know
			// that the original was successfully output in upper case
			if strings.ToLower(s) != s {
				return nil
			}
			return errors.New("Country was returned in lowercase, but was requested in uppercase")
		},
	},
	{
		Template: "{country:case:down}",
		Comparator: func(s string) error {
			// Since I can't know which country comes out, i'll invert the result
			// If the ToLowered result is not the same as the original result, we know
			// that the original was successfully output in upper case
			if strings.ToUpper(s) != s {
				return nil
			}
			return errors.New("Country was returned in uppercase, but was requested in lowercase")
		},
	},
	{
		Template: "{country}@{country:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Country at position 1 not equal to country at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{country}@{country:ordinal:1}",
		WriteFailure: true,
	},
}

// Placeholders
var FloatCases = []TestCase{
	{
		Template: "{float}",
		Comparator: func(s string) error {
			i, err := strconv.ParseFloat(s, 64)
			if err != nil {
				return err
			}
			if i >= 0.0 && i <= 100.0 {
				return nil
			}
			return errors.New("Float out of range for default min/max values")
		},
	},
	{
		Template: "{float:max:5000.0|min:4999.0}",
		Comparator: func(s string) error {
			i, err := strconv.ParseFloat(s, 64)
			if err != nil {
				return err
			}
			if i >= 4999.0 && i <= 5000.0 {
				return nil
			}
			return errors.New("Float out of range for custom min/max values")
		},
	},
	{
		Template: "{float:max:-5000.0|min:-5001.0}",
		Comparator: func(s string) error {
			i, err := strconv.ParseFloat(s, 64)
			if err != nil {
				return err
			}
			if i >= -5001.0 && i <= -5000.0 {
				return nil
			}
			return errors.New("Float out of range for custom min/max values")
		},
	},
	{
		Template: "{float}@{float:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Float at position 1 not equal to Float at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{float}@{float:ordinal:1}",
		WriteFailure: true,
	},
}

var IntegerCases = []TestCase{
	{
		Template: "{int}",
		Comparator: func(s string) error {
			i, err := strconv.Atoi(s)
			if err != nil {
				return err
			}
			if i >= 0 && i <= 100 {
				return nil
			}
			return errors.New("Int out of range for default min/max values")
		},
	},
	{
		Template: "{int:max:5000|min:4999}",
		Comparator: func(s string) error {
			i, err := strconv.Atoi(s)
			if err != nil {
				return err
			}
			if i >= 4999 && i <= 5000 {
				return nil
			}
			return errors.New("Int out of range for custom min/max values")
		},
	},
	{
		Template: "{int:max:-5000|min:-5001}",
		Comparator: func(s string) error {
			i, err := strconv.Atoi(s)
			if err != nil {
				return err
			}
			if i >= -5001 && i <= -5000 {
				return nil
			}
			return errors.New("Int out of range for custom min/max values")
		},
	},
	{
		Template: "{int}@{int:ordinal:0}",
		Comparator: func(s string) error {
			p := strings.Split(s, "@")
			if p[0] == p[1] {
				return nil
			}
			return errors.New("Int at position 1 not equal to int at position 0: " + p[0] + " " + p[1])
		},
	},
	{
		Template:     "{int}@{int:ordinal:1}",
		WriteFailure: true,
	},
}

var AllCases = [][]TestCase{
	GUIDCases,
	NowCases,
	TimeCases,
	CountryCases,
	FloatCases,
	IntegerCases,
}

// TODO Test each random function individually, under a number of inputs to make supported
// all the options behave as expected.

func TestMain(m *testing.M) {
	rand.Seed(time.Now().Unix())
	os.Exit(m.Run())
}

func TestAllCases(t *testing.T) {
	// TODO The library should be threadsafe, I should go wide here to run all specs
	// in parallel, like the natural tests would be. Channel + waitgroup to collect
	// and report on errors once they all finish
	for _, cs := range AllCases {
		for _, c := range cs {
			cs, err := BuildCallstack(c.Template)
			// If we get an error and weren't expecting it
			// Or, if we didn't get one but were expecting it
			if err != nil && !c.ParseFailure {
				t.Error(err)
			} else if err == nil && c.ParseFailure {
				t.Error("Expected to encounter Parse Failure, but did not for Test Case ", c.Template)
			}

			result := &bytes.Buffer{}
			err = cs.Write(result)

			// If we get an error and weren't expecting it
			// Or, if we didn't get one but were expecting it
			if err != nil && !c.WriteFailure {
				t.Error(err)
			} else if err == nil && c.ParseFailure {
				t.Error("Expected to encounter Write Failure, but did not for Test Case ", c.Template)
			}

			if c.Comparator != nil {
				if err := c.Comparator(result.String()); err != nil {
					t.Error(err)
				}
			}
		}
	}
}

func BenchmarkWrites(b *testing.B) {
	template := "INSERT INTO floof VALUES ('{guid}','{time},'{guid:ordinal:0}','{country}',{int:min:-2000|max:0},{int:min:100|max:1000},{float:min:-1000.0|max:-540.0},{int:min:1|max:40},'{now}','{now:ordinal:0}','{unicode:length:2|case:up}',NULL,-3)"
	var cs *Callstack
	var err error
	if cs, err = BuildCallstack(template); err != nil {
		b.Error(err)
	}
	for n := 0; n < b.N; n++ {
		result := &bytes.Buffer{}
		err = cs.Write(result)
		if err != nil {
			b.Error(err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import "fmt"

type stringer interface {
	String() string
}

func concat[T stringer](vals []T) string {
	result := ""
	for _, val := range vals {
		result += val.String()
	}
	return result
}

type car struct {
	price int
}

func (c car) String() string {
	return fmt.Sprintf("%d", c.price)
}

func main() {
	val := concat([]stringer{
		car{price: 1},
		car{price: 2},
	})

	fmt.Println(val)
}
<commit_msg>chore(example46): custom interface constrain<commit_after>package main

import "fmt"

type stringer interface {
	String() string
}

func concat[T stringer](vals []T) string {
	result := ""
	for _, val := range vals {
		result += val.String()
	}
	return result
}

type car struct {
	price int
}

func (c car) String() string {
	return fmt.Sprintf("%d", c.price)
}

type plusser interface {
	Plus(string) string
}

func cconcatTo[S stringer, P plusser](s []S, p []P) []string {
	r := make([]string, len(s))
	for i, v := range s {
		r[i] = p[i].Plus(v.String())
	}
	return r
}

type foo struct {
	name string
}

func (c foo) String() string {
	return c.name
}

type plus map[string]string

func (p plus) Plus(k string) string {
	if v, ok := p[k]; ok {
		return v
	}
	return ""
}

func main() {
	// exmaple01
	val := concat([]stringer{
		car{price: 1},
		car{price: 2},
	})

	fmt.Println(val)

	// example02
	p := plus{
		"a": "100",
		"b": "200",
	}

	fmt.Println(cconcatTo([]stringer{
		foo{name: "a"},
		foo{name: "b"},
	}, []plusser{p, p}))
}
<|endoftext|>