<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"syscall"
	"time"

	"github.com/mercari/gaurun/gaurun"
)

const (
	DefaultPidPermission = 0644
)

func main() {
	versionPrinted := flag.Bool("v", false, "gaurun version")
	confPath := flag.String("c", "", "configuration file path for gaurun")
	listenPort := flag.String("p", "", "port number or unix socket path")
	workerNum := flag.Int64("w", 0, "number of workers for push notification")
	queueNum := flag.Int64("q", 0, "size of internal queue for push notification")
	flag.Parse()

	if *versionPrinted {
		gaurun.PrintVersion()
		return
	}

	// set default parameters
	gaurun.ConfGaurun = gaurun.BuildDefaultConf()

	// load configuration
	conf, err := gaurun.LoadConf(gaurun.ConfGaurun, *confPath)
	if err != nil {
		gaurun.LogSetupFatal(err)
	}
	gaurun.ConfGaurun = conf

	// overwrite if port is specified by flags
	if *listenPort != "" {
		gaurun.ConfGaurun.Core.Port = *listenPort
	}

	// overwrite if workerNum is specified by flags
	if *workerNum > 0 {
		gaurun.ConfGaurun.Core.WorkerNum = *workerNum
	}

	// overwrite if queueNum is specified by flags
	if *queueNum > 0 {
		gaurun.ConfGaurun.Core.QueueNum = *queueNum
	}

	// set logger
	accessLogger, accessLogReopener, err := gaurun.InitLog(gaurun.ConfGaurun.Log.AccessLog, "info")
	if err != nil {
		gaurun.LogSetupFatal(err)
	}
	errorLogger, errorLogReopener, err := gaurun.InitLog(gaurun.ConfGaurun.Log.ErrorLog, gaurun.ConfGaurun.Log.Level)
	if err != nil {
		gaurun.LogSetupFatal(err)
	}

	gaurun.LogAccess = accessLogger
	gaurun.LogError = errorLogger

	if !gaurun.ConfGaurun.Ios.Enabled && !gaurun.ConfGaurun.Android.Enabled {
		gaurun.LogSetupFatal(fmt.Errorf("What do you want to do?"))
	}

	if gaurun.ConfGaurun.Ios.Enabled {
		gaurun.CertificatePemIos.Cert, err = ioutil.ReadFile(gaurun.ConfGaurun.Ios.PemCertPath)
		if err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("A certification file for iOS is not found."))
		}

		gaurun.CertificatePemIos.Key, err = ioutil.ReadFile(gaurun.ConfGaurun.Ios.PemKeyPath)
		if err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("A key file for iOS is not found."))
		}

	}

	if gaurun.ConfGaurun.Android.Enabled {
		if gaurun.ConfGaurun.Android.ApiKey == "" {
			gaurun.LogSetupFatal(fmt.Errorf("APIKey for Android is empty."))
		}
	}

	sigHUPChan := make(chan os.Signal, 1)
	signal.Notify(sigHUPChan, syscall.SIGHUP)

	sighupHandler := func() {
		if err := accessLogReopener.Reopen(); err != nil {
			gaurun.LogError.Warn(fmt.Sprintf("failed to reopen access log: %v", err))
		}
		if err := errorLogReopener.Reopen(); err != nil {
			gaurun.LogError.Warn(fmt.Sprintf("failed to reopen error log: %v", err))
		}
	}

	go signalHandler(sigHUPChan, sighupHandler)

	if len(conf.Core.Pid) > 0 {
		if _, err := os.Stat(filepath.Dir(conf.Core.Pid)); os.IsNotExist(err) {
			gaurun.LogSetupFatal(fmt.Errorf("directory for pid file is not exist: %v", err))
		} else if err := ioutil.WriteFile(conf.Core.Pid, []byte(strconv.Itoa(os.Getpid())), DefaultPidPermission); err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("failed to create a pid file: %v", err))
		}
	}

	if gaurun.ConfGaurun.Android.Enabled {
		if err := gaurun.InitGCMClient(); err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("failed to init gcm/fcm client: %v", err))
		}
	}

	if gaurun.ConfGaurun.Ios.Enabled {
		if err := gaurun.InitAPNSClient(); err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("failed to init http client for APNs: %v", err))
		}
	}
	gaurun.InitStat()
	gaurun.StartPushWorkers(gaurun.ConfGaurun.Core.WorkerNum, gaurun.ConfGaurun.Core.QueueNum)

	mux := http.NewServeMux()
	gaurun.RegisterHandlers(mux)

	server := &http.Server{
		Handler: mux,
	}
	go func() {
		gaurun.LogError.Info("start server")
		if err := gaurun.RunServer(server, &gaurun.ConfGaurun); err != nil {
			gaurun.LogError.Info(fmt.Sprintf("failed to serve: %s", err))
		}
	}()

	// Graceful shutdown (kicked by SIGTERM).
	//
	// First, it shutdowns server and stops accepting new requests.
	// Then wait until all remaining queues in buffer are flushed.
	sigTERMChan := make(chan os.Signal, 1)
	signal.Notify(sigTERMChan, syscall.SIGTERM)

	<-sigTERMChan
	gaurun.LogError.Info("shutdown server")
	timeout := time.Duration(conf.Core.ShutdownTimeout) * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	if err := server.Shutdown(ctx); err != nil {
		gaurun.LogError.Error(fmt.Sprintf("failed to shutdown server: %v", err))
	}

	// Start a goroutine to log number of job queue.
	go func() {
		for {
			queue := len(gaurun.QueueNotification)
			if queue == 0 {
				break
			}

			gaurun.LogError.Info(fmt.Sprintf("wait until queue is empty. Current queue len: %d", queue))
			time.Sleep(1 * time.Second)
		}
	}()

	// Block until all pusher worker job is done.
	gaurun.PusherWg.Wait()

	gaurun.LogError.Info("successfully shutdown")
}

func signalHandler(ch <-chan os.Signal, sighupFn func()) {
	for {
		select {
		case sig := <-ch:
			switch sig {
			case syscall.SIGHUP:
				sighupFn()
			}
		}
	}
}
<commit_msg>fix error strings format<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"syscall"
	"time"

	"github.com/mercari/gaurun/gaurun"
)

const (
	DefaultPidPermission = 0644
)

func main() {
	versionPrinted := flag.Bool("v", false, "gaurun version")
	confPath := flag.String("c", "", "configuration file path for gaurun")
	listenPort := flag.String("p", "", "port number or unix socket path")
	workerNum := flag.Int64("w", 0, "number of workers for push notification")
	queueNum := flag.Int64("q", 0, "size of internal queue for push notification")
	flag.Parse()

	if *versionPrinted {
		gaurun.PrintVersion()
		return
	}

	// set default parameters
	gaurun.ConfGaurun = gaurun.BuildDefaultConf()

	// load configuration
	conf, err := gaurun.LoadConf(gaurun.ConfGaurun, *confPath)
	if err != nil {
		gaurun.LogSetupFatal(err)
	}
	gaurun.ConfGaurun = conf

	// overwrite if port is specified by flags
	if *listenPort != "" {
		gaurun.ConfGaurun.Core.Port = *listenPort
	}

	// overwrite if workerNum is specified by flags
	if *workerNum > 0 {
		gaurun.ConfGaurun.Core.WorkerNum = *workerNum
	}

	// overwrite if queueNum is specified by flags
	if *queueNum > 0 {
		gaurun.ConfGaurun.Core.QueueNum = *queueNum
	}

	// set logger
	accessLogger, accessLogReopener, err := gaurun.InitLog(gaurun.ConfGaurun.Log.AccessLog, "info")
	if err != nil {
		gaurun.LogSetupFatal(err)
	}
	errorLogger, errorLogReopener, err := gaurun.InitLog(gaurun.ConfGaurun.Log.ErrorLog, gaurun.ConfGaurun.Log.Level)
	if err != nil {
		gaurun.LogSetupFatal(err)
	}

	gaurun.LogAccess = accessLogger
	gaurun.LogError = errorLogger

	if !gaurun.ConfGaurun.Ios.Enabled && !gaurun.ConfGaurun.Android.Enabled {
		gaurun.LogSetupFatal(fmt.Errorf("no platform has been enabled"))
	}

	if gaurun.ConfGaurun.Ios.Enabled {
		gaurun.CertificatePemIos.Cert, err = ioutil.ReadFile(gaurun.ConfGaurun.Ios.PemCertPath)
		if err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("the certification file for iOS was not found"))
		}

		gaurun.CertificatePemIos.Key, err = ioutil.ReadFile(gaurun.ConfGaurun.Ios.PemKeyPath)
		if err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("the key file for iOS was not found"))
		}

	}

	if gaurun.ConfGaurun.Android.Enabled {
		if gaurun.ConfGaurun.Android.ApiKey == "" {
			gaurun.LogSetupFatal(fmt.Errorf("the APIKey for Android cannot be empty"))
		}
	}

	sigHUPChan := make(chan os.Signal, 1)
	signal.Notify(sigHUPChan, syscall.SIGHUP)

	sighupHandler := func() {
		if err := accessLogReopener.Reopen(); err != nil {
			gaurun.LogError.Warn(fmt.Sprintf("failed to reopen access log: %v", err))
		}
		if err := errorLogReopener.Reopen(); err != nil {
			gaurun.LogError.Warn(fmt.Sprintf("failed to reopen error log: %v", err))
		}
	}

	go signalHandler(sigHUPChan, sighupHandler)

	if len(conf.Core.Pid) > 0 {
		if _, err := os.Stat(filepath.Dir(conf.Core.Pid)); os.IsNotExist(err) {
			gaurun.LogSetupFatal(fmt.Errorf("directory for pid file is not exist: %v", err))
		} else if err := ioutil.WriteFile(conf.Core.Pid, []byte(strconv.Itoa(os.Getpid())), DefaultPidPermission); err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("failed to create a pid file: %v", err))
		}
	}

	if gaurun.ConfGaurun.Android.Enabled {
		if err := gaurun.InitGCMClient(); err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("failed to init gcm/fcm client: %v", err))
		}
	}

	if gaurun.ConfGaurun.Ios.Enabled {
		if err := gaurun.InitAPNSClient(); err != nil {
			gaurun.LogSetupFatal(fmt.Errorf("failed to init http client for APNs: %v", err))
		}
	}
	gaurun.InitStat()
	gaurun.StartPushWorkers(gaurun.ConfGaurun.Core.WorkerNum, gaurun.ConfGaurun.Core.QueueNum)

	mux := http.NewServeMux()
	gaurun.RegisterHandlers(mux)

	server := &http.Server{
		Handler: mux,
	}
	go func() {
		gaurun.LogError.Info("start server")
		if err := gaurun.RunServer(server, &gaurun.ConfGaurun); err != nil {
			gaurun.LogError.Info(fmt.Sprintf("failed to serve: %s", err))
		}
	}()

	// Graceful shutdown (kicked by SIGTERM).
	//
	// First, it shutdowns server and stops accepting new requests.
	// Then wait until all remaining queues in buffer are flushed.
	sigTERMChan := make(chan os.Signal, 1)
	signal.Notify(sigTERMChan, syscall.SIGTERM)

	<-sigTERMChan
	gaurun.LogError.Info("shutdown server")
	timeout := time.Duration(conf.Core.ShutdownTimeout) * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	if err := server.Shutdown(ctx); err != nil {
		gaurun.LogError.Error(fmt.Sprintf("failed to shutdown server: %v", err))
	}

	// Start a goroutine to log number of job queue.
	go func() {
		for {
			queue := len(gaurun.QueueNotification)
			if queue == 0 {
				break
			}

			gaurun.LogError.Info(fmt.Sprintf("wait until queue is empty. Current queue len: %d", queue))
			time.Sleep(1 * time.Second)
		}
	}()

	// Block until all pusher worker job is done.
	gaurun.PusherWg.Wait()

	gaurun.LogError.Info("successfully shutdown")
}

func signalHandler(ch <-chan os.Signal, sighupFn func()) {
	for {
		select {
		case sig := <-ch:
			switch sig {
			case syscall.SIGHUP:
				sighupFn()
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/sqs/go-deps"
	"os"
)

func main() {
	flag.Parse()
	if flag.NArg() != 1 {
		fmt.Fprintf(os.Stderr, "must specify a single, locally installed package\n")
		os.Exit(1)
	}

	importPath := flag.Arg(0)
	pkg, err := deps.Read(importPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error reading package: %s\n", err)
		os.Exit(1)
	}

	if pkg.Dir == "" {
		// User specified package that's not installed locally.
		fmt.Fprintf(os.Stderr, "can't find package: import path %s\n", importPath)
		os.Exit(1)
	}

	fmt.Printf("# %s\n", importPath)
	for _, p := range pkg.DepsNotFound {
		fmt.Printf("%s\n", p)
		if err := deps.Download(p); err != nil {
			fmt.Fprintf(os.Stderr, "can't download dep: import path %s: %s\n", p, err)
			os.Exit(1)
		}
	}
}
<commit_msg>Remove unnecessary godeps cmd<commit_after><|endoftext|>


<commit_before>package main

import (
	"log"
	"os"

	"github.com/aybabtme/humanlog"
	"github.com/aybabtme/rgbterm"
	"github.com/urfave/cli"
)

var version = "devel"

func fatalf(c *cli.Context, format string, args ...interface{}) {
	log.Printf(format, args...)
	cli.ShowAppHelp(c)
	os.Exit(1)
}

func main() {
	app := newApp()

	prefix := rgbterm.FgString(app.Name+"> ", 99, 99, 99)

	log.SetFlags(0)
	log.SetPrefix(prefix)
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func newApp() *cli.App {

	skip := cli.StringSlice{}
	keep := cli.StringSlice{}

	skipFlag := cli.StringSliceFlag{
		Name:  "skip",
		Usage: "keys to skip when parsing a log entry",
		Value: &skip,
	}

	keepFlag := cli.StringSliceFlag{
		Name:  "keep",
		Usage: "keys to keep when parsing a log entry",
		Value: &keep,
	}

	sortLongest := cli.BoolTFlag{
		Name:  "sort-longest",
		Usage: "sort by longest key after having sorted lexicographically",
	}

	skipUnchanged := cli.BoolTFlag{
		Name:  "skip-unchanged",
		Usage: "skip keys that have the same value than the previous entry",
	}

	truncates := cli.BoolFlag{
		Name:  "truncate",
		Usage: "truncates values that are longer than --truncate-length",
	}

	truncateLength := cli.IntFlag{
		Name:  "truncate-length",
		Usage: "truncate values that are longer than this length",
		Value: humanlog.DefaultOptions.TruncateLength,
	}

	app := cli.NewApp()
	app.Author = "Antoine Grondin"
	app.Email = "antoine@digitalocean.com"
	app.Name = "humanlog"
	app.Version = version
	app.Usage = "reads structured logs from stdin, makes them pretty on stdout!"

	app.Flags = []cli.Flag{skipFlag, keepFlag, sortLongest, skipUnchanged, truncates, truncateLength}

	app.Action = func(c *cli.Context) {

		opts := humanlog.DefaultOptions
		opts.SortLongest = c.BoolT(sortLongest.Name)
		opts.SkipUnchanged = c.BoolT(skipUnchanged.Name)
		opts.Truncates = c.BoolT(truncates.Name)
		opts.TruncateLength = c.Int(truncateLength.Name)

		switch {
		case c.IsSet(skipFlag.Name) && c.IsSet(keepFlag.Name):
			fatalf(c, "can only use one of %q and %q", skipFlag.Name, keepFlag.Name)
		case c.IsSet(skipFlag.Name):
			opts.SetSkip(skip)
		case c.IsSet(keepFlag.Name):
			opts.SetKeep(keep)
		}

		log.Print("reading stdin...")
		if err := humanlog.Scanner(os.Stdin, os.Stdout, opts); err != nil {
			log.Fatalf("scanning caught an error: %v", err)
		}
	}
	return app
}
<commit_msg>update for cli.urfave<commit_after>package main

import (
	"log"
	"os"

	"github.com/aybabtme/humanlog"
	"github.com/aybabtme/rgbterm"
	"github.com/urfave/cli"
)

var version = "devel"

func fatalf(c *cli.Context, format string, args ...interface{}) {
	log.Printf(format, args...)
	cli.ShowAppHelp(c)
	os.Exit(1)
}

func main() {
	app := newApp()

	prefix := rgbterm.FgString(app.Name+"> ", 99, 99, 99)

	log.SetFlags(0)
	log.SetPrefix(prefix)
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func newApp() *cli.App {

	skip := cli.StringSlice{}
	keep := cli.StringSlice{}

	skipFlag := cli.StringSliceFlag{
		Name:  "skip",
		Usage: "keys to skip when parsing a log entry",
		Value: &skip,
	}

	keepFlag := cli.StringSliceFlag{
		Name:  "keep",
		Usage: "keys to keep when parsing a log entry",
		Value: &keep,
	}

	sortLongest := cli.BoolTFlag{
		Name:  "sort-longest",
		Usage: "sort by longest key after having sorted lexicographically",
	}

	skipUnchanged := cli.BoolTFlag{
		Name:  "skip-unchanged",
		Usage: "skip keys that have the same value than the previous entry",
	}

	truncates := cli.BoolFlag{
		Name:  "truncate",
		Usage: "truncates values that are longer than --truncate-length",
	}

	truncateLength := cli.IntFlag{
		Name:  "truncate-length",
		Usage: "truncate values that are longer than this length",
		Value: humanlog.DefaultOptions.TruncateLength,
	}

	app := cli.NewApp()
	app.Author = "Antoine Grondin"
	app.Email = "antoine@digitalocean.com"
	app.Name = "humanlog"
	app.Version = version
	app.Usage = "reads structured logs from stdin, makes them pretty on stdout!"

	app.Flags = []cli.Flag{skipFlag, keepFlag, sortLongest, skipUnchanged, truncates, truncateLength}

	app.Action = func(c *cli.Context) error {

		opts := humanlog.DefaultOptions
		opts.SortLongest = c.BoolT(sortLongest.Name)
		opts.SkipUnchanged = c.BoolT(skipUnchanged.Name)
		opts.Truncates = c.BoolT(truncates.Name)
		opts.TruncateLength = c.Int(truncateLength.Name)

		switch {
		case c.IsSet(skipFlag.Name) && c.IsSet(keepFlag.Name):
			fatalf(c, "can only use one of %q and %q", skipFlag.Name, keepFlag.Name)
		case c.IsSet(skipFlag.Name):
			opts.SetSkip(skip)
		case c.IsSet(keepFlag.Name):
			opts.SetKeep(keep)
		}

		log.Print("reading stdin...")
		if err := humanlog.Scanner(os.Stdin, os.Stdout, opts); err != nil {
			log.Fatalf("scanning caught an error: %v", err)
		}
		return nil
	}
	return app
}
<|endoftext|>


<commit_before>/*
i-sudoku is an interactive command-line sudoku tool
*/

package main

import (
	"github.com/jkomoros/sudoku"
	"github.com/nsf/termbox-go"
	"log"
	"strconv"
	"strings"
	"unicode/utf8"
)

type mainModel struct {
	grid         *sudoku.Grid
	Selected     *sudoku.Cell
	marksToInput []int
}

func main() {
	if err := termbox.Init(); err != nil {
		log.Fatal("Termbox initialization failed:", err)
	}
	defer termbox.Close()

	model := newModel()

	draw(model)

mainloop:
	for {
		switch ev := termbox.PollEvent(); ev.Type {
		case termbox.EventKey:
			switch ev.Key {
			case termbox.KeyEsc:
				if model.ModeInputEsc() {
					break mainloop
				}
			case termbox.KeyCtrlC:
				break mainloop
			case termbox.KeyArrowDown:
				model.MoveSelectionDown()
			case termbox.KeyArrowLeft:
				model.MoveSelectionLeft()
			case termbox.KeyArrowRight:
				model.MoveSelectionRight()
			case termbox.KeyArrowUp:
				model.MoveSelectionUp()
			case termbox.KeyEnter:
				model.ModeCommitMarkMode()
			}
			switch ev.Ch {
			case 'q':
				break mainloop
			case 'm':
				model.ModeEnterMarkMode()
			case 'n':
				//TODO: since this is a destructive action, require a confirmation
				model.NewGrid()
			//TODO: do this in a more general way related to DIM
			case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
				//TODO: this is a seriously gross way of converting a rune to a string.
				num, err := strconv.Atoi(strings.Replace(strconv.QuoteRuneToASCII(ev.Ch), "'", "", -1))
				if err != nil {
					panic(err)
				}
				model.ModeInputNumber(num)
			}
		}
		draw(model)
	}
}

func newModel() *mainModel {
	model := &mainModel{}
	model.EnsureSelected()
	return model
}

func (m *mainModel) ModeInputEsc() (quit bool) {
	if m.marksToInput != nil {
		m.ModeCancelMarkMode()
		return false
	}
	return true
}

func (m *mainModel) ModeEnterMarkMode() {
	m.marksToInput = make([]int, 0)
}

func (m *mainModel) ModeCommitMarkMode() {
	for _, num := range m.marksToInput {
		m.ToggleSelectedMark(num)
	}
	m.marksToInput = nil
}

func (m *mainModel) ModeCancelMarkMode() {
	m.marksToInput = nil
}

func (m *mainModel) ModeInputNumber(num int) {
	if m.marksToInput == nil {
		m.SetSelectedNumber(num)
	} else {
		m.marksToInput = append(m.marksToInput, num)
	}
}

func (m *mainModel) EnsureSelected() {
	m.EnsureGrid()
	//Ensures that at least one cell is selected.
	if m.Selected == nil {
		m.Selected = m.grid.Cell(0, 0)
	}
}

func (m *mainModel) MoveSelectionLeft() {
	m.EnsureSelected()
	r := m.Selected.Row()
	c := m.Selected.Col()
	c--
	if c < 0 {
		c = 0
	}
	m.Selected = m.grid.Cell(r, c)
}

func (m *mainModel) MoveSelectionRight() {
	m.EnsureSelected()
	r := m.Selected.Row()
	c := m.Selected.Col()
	c++
	if c >= sudoku.DIM {
		c = sudoku.DIM - 1
	}
	m.Selected = m.grid.Cell(r, c)
}

func (m *mainModel) MoveSelectionUp() {
	m.EnsureSelected()
	r := m.Selected.Row()
	c := m.Selected.Col()
	r--
	if r < 0 {
		r = 0
	}
	m.Selected = m.grid.Cell(r, c)
}

func (m *mainModel) MoveSelectionDown() {
	m.EnsureSelected()
	r := m.Selected.Row()
	c := m.Selected.Col()
	r++
	if r >= sudoku.DIM {
		r = sudoku.DIM - 1
	}
	m.Selected = m.grid.Cell(r, c)
}

func (m *mainModel) EnsureGrid() {
	if m.grid == nil {
		m.NewGrid()
	}
}

func (m *mainModel) NewGrid() {
	m.grid = sudoku.GenerateGrid(nil)
	m.grid.LockFilledCells()
}

func (m *mainModel) SetSelectedNumber(num int) {
	m.EnsureSelected()
	if m.Selected.Locked() {
		return
	}
	m.Selected.SetNumber(num)
}

func (m *mainModel) ToggleSelectedMark(num int) {
	m.EnsureSelected()
	if m.Selected.Locked() {
		return
	}
	m.Selected.SetMark(num, !m.Selected.Mark(num))
}

func draw(model *mainModel) {
	drawGrid(model)
	termbox.Flush()
}

func drawGrid(model *mainModel) {

	//TODO: have a mode line after the grid for if the grid is invalid, if it's solved.

	//TODO: draw which marks are ready to commit.

	grid := model.grid

	selectedTop, selectedLeft, selectedHeight, selectedWidth := model.Selected.DiagramExtents()

	for y, line := range strings.Split(grid.Diagram(true), "\n") {
		x := 0
		//The first number in range will be byte offset, but for some items like the bullet, it's two bytes.
		//But what we care about is that each item is a character.
		for _, ch := range line {

			defaultColor := termbox.ColorGreen

			numberRune, _ := utf8.DecodeRuneInString(sudoku.DIAGRAM_NUMBER)
			lockedRune, _ := utf8.DecodeRuneInString(sudoku.DIAGRAM_LOCKED)

			if ch == numberRune {
				defaultColor = termbox.ColorBlue
			} else if ch == lockedRune {
				defaultColor = termbox.ColorRed
			}

			backgroundColor := termbox.ColorDefault

			if x >= selectedTop && x < (selectedTop+selectedHeight) && y >= selectedLeft && y < (selectedLeft+selectedWidth) {
				//We're on the selected cell
				backgroundColor = termbox.ColorWhite
			}

			termbox.SetCell(x, y, ch, defaultColor, backgroundColor)
			x++
		}
	}
}
<commit_msg>Added a TODO<commit_after>/*
i-sudoku is an interactive command-line sudoku tool
*/

package main

import (
	"github.com/jkomoros/sudoku"
	"github.com/nsf/termbox-go"
	"log"
	"strconv"
	"strings"
	"unicode/utf8"
)

type mainModel struct {
	grid         *sudoku.Grid
	Selected     *sudoku.Cell
	marksToInput []int
}

func main() {
	if err := termbox.Init(); err != nil {
		log.Fatal("Termbox initialization failed:", err)
	}
	defer termbox.Close()

	model := newModel()

	draw(model)

mainloop:
	for {
		switch ev := termbox.PollEvent(); ev.Type {
		case termbox.EventKey:
			switch ev.Key {
			case termbox.KeyEsc:
				if model.ModeInputEsc() {
					break mainloop
				}
			case termbox.KeyCtrlC:
				break mainloop
			case termbox.KeyArrowDown:
				model.MoveSelectionDown()
			case termbox.KeyArrowLeft:
				model.MoveSelectionLeft()
			case termbox.KeyArrowRight:
				model.MoveSelectionRight()
			case termbox.KeyArrowUp:
				model.MoveSelectionUp()
			case termbox.KeyEnter:
				model.ModeCommitMarkMode()
			}
			switch ev.Ch {
			case 'q':
				break mainloop
			case 'm':
				//TODO: ideally Ctrl+Num would work to put in one mark. But termbox doesn't appear to let that work.
				model.ModeEnterMarkMode()
			case 'n':
				//TODO: since this is a destructive action, require a confirmation
				model.NewGrid()
			//TODO: do this in a more general way related to DIM
			case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
				//TODO: this is a seriously gross way of converting a rune to a string.
				num, err := strconv.Atoi(strings.Replace(strconv.QuoteRuneToASCII(ev.Ch), "'", "", -1))
				if err != nil {
					panic(err)
				}
				model.ModeInputNumber(num)
			}
		}
		draw(model)
	}
}

func newModel() *mainModel {
	model := &mainModel{}
	model.EnsureSelected()
	return model
}

func (m *mainModel) ModeInputEsc() (quit bool) {
	if m.marksToInput != nil {
		m.ModeCancelMarkMode()
		return false
	}
	return true
}

func (m *mainModel) ModeEnterMarkMode() {
	m.marksToInput = make([]int, 0)
}

func (m *mainModel) ModeCommitMarkMode() {
	for _, num := range m.marksToInput {
		m.ToggleSelectedMark(num)
	}
	m.marksToInput = nil
}

func (m *mainModel) ModeCancelMarkMode() {
	m.marksToInput = nil
}

func (m *mainModel) ModeInputNumber(num int) {
	if m.marksToInput == nil {
		m.SetSelectedNumber(num)
	} else {
		m.marksToInput = append(m.marksToInput, num)
	}
}

func (m *mainModel) EnsureSelected() {
	m.EnsureGrid()
	//Ensures that at least one cell is selected.
	if m.Selected == nil {
		m.Selected = m.grid.Cell(0, 0)
	}
}

func (m *mainModel) MoveSelectionLeft() {
	m.EnsureSelected()
	r := m.Selected.Row()
	c := m.Selected.Col()
	c--
	if c < 0 {
		c = 0
	}
	m.Selected = m.grid.Cell(r, c)
}

func (m *mainModel) MoveSelectionRight() {
	m.EnsureSelected()
	r := m.Selected.Row()
	c := m.Selected.Col()
	c++
	if c >= sudoku.DIM {
		c = sudoku.DIM - 1
	}
	m.Selected = m.grid.Cell(r, c)
}

func (m *mainModel) MoveSelectionUp() {
	m.EnsureSelected()
	r := m.Selected.Row()
	c := m.Selected.Col()
	r--
	if r < 0 {
		r = 0
	}
	m.Selected = m.grid.Cell(r, c)
}

func (m *mainModel) MoveSelectionDown() {
	m.EnsureSelected()
	r := m.Selected.Row()
	c := m.Selected.Col()
	r++
	if r >= sudoku.DIM {
		r = sudoku.DIM - 1
	}
	m.Selected = m.grid.Cell(r, c)
}

func (m *mainModel) EnsureGrid() {
	if m.grid == nil {
		m.NewGrid()
	}
}

func (m *mainModel) NewGrid() {
	m.grid = sudoku.GenerateGrid(nil)
	m.grid.LockFilledCells()
}

func (m *mainModel) SetSelectedNumber(num int) {
	m.EnsureSelected()
	if m.Selected.Locked() {
		return
	}
	m.Selected.SetNumber(num)
}

func (m *mainModel) ToggleSelectedMark(num int) {
	m.EnsureSelected()
	if m.Selected.Locked() {
		return
	}
	m.Selected.SetMark(num, !m.Selected.Mark(num))
}

func draw(model *mainModel) {
	drawGrid(model)
	termbox.Flush()
}

func drawGrid(model *mainModel) {

	//TODO: have a mode line after the grid for if the grid is invalid, if it's solved.

	//TODO: draw which marks are ready to commit.

	grid := model.grid

	selectedTop, selectedLeft, selectedHeight, selectedWidth := model.Selected.DiagramExtents()

	for y, line := range strings.Split(grid.Diagram(true), "\n") {
		x := 0
		//The first number in range will be byte offset, but for some items like the bullet, it's two bytes.
		//But what we care about is that each item is a character.
		for _, ch := range line {

			defaultColor := termbox.ColorGreen

			numberRune, _ := utf8.DecodeRuneInString(sudoku.DIAGRAM_NUMBER)
			lockedRune, _ := utf8.DecodeRuneInString(sudoku.DIAGRAM_LOCKED)

			if ch == numberRune {
				defaultColor = termbox.ColorBlue
			} else if ch == lockedRune {
				defaultColor = termbox.ColorRed
			}

			backgroundColor := termbox.ColorDefault

			if x >= selectedTop && x < (selectedTop+selectedHeight) && y >= selectedLeft && y < (selectedLeft+selectedWidth) {
				//We're on the selected cell
				backgroundColor = termbox.ColorWhite
			}

			termbox.SetCell(x, y, ch, defaultColor, backgroundColor)
			x++
		}
	}
}
<|endoftext|>


<commit_before>// Package main defines a command line tool for submitting date
// ranges for reprocessing
package main

// TODO - note about setting up batch table and giving permission.

/*
Strategies...
  1. Work from a prefix, or range of prefixes.
  2. Work from a date range
  3. Work from a month prefix, but explicitly iterate over days.
      maybe use a separate goroutine for each date?

Usage:

*/

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"sync"
	"sync/atomic"

	"github.com/spaolacci/murmur3"
	"google.golang.org/api/iterator"

	"cloud.google.com/go/storage"
)

var (
	fProject   = flag.String("project", "mlab-oti", "Project containing queues.")
	fQueue     = flag.String("queue", "etl-ndt-batch-", "Base of queue name.")
	fNumQueues = flag.Int("num_queues", 5, "Number of queues.  Normally determined by listing queues.")
	fBucket    = flag.String("bucket", "archive-mlab-oti", "Source bucket.")
	fExper     = flag.String("experiment", "ndt", "Experiment prefix, trailing slash optional")
	fMonth     = flag.String("month", "", "Single month spec, as YYYY/MM")
	fDay       = flag.String("day", "", "Single day spec, as YYYY/MM/DD")

	errCount      int32
	storageClient *storage.Client
	bucket        *storage.BucketHandle

	hasher = murmur3.New32()
)

func init() {
	// Always prepend the filename and line number.
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func postOne(queue string, bucket string, fn string) error {
	reqStr := fmt.Sprintf("http://queue-pusher-dot-%s.appspot.com/receiver?queue=%s&filename=gs://%s/%s", *fProject, queue, bucket, fn)
	resp, err := http.Get(reqStr)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return errors.New("http error: " + resp.Status)
	}

	return nil
}

// Post all items in an ObjectIterator into specific
// queue.
func postDay(wg *sync.WaitGroup, queue string, it *storage.ObjectIterator) {
	defer wg.Done()
	log.Printf("%+v\n", it)
	for o, err := it.Next(); err != iterator.Done; o, err = it.Next() {
		if err != nil {
			log.Println(err)
			ec := atomic.AddInt32(&errCount, 1)
			if ec > 10 {
				panic(err)
			}
		}

		err = postOne(queue, *fBucket, o.Name)
		if err != nil {
			log.Println(err)
			ec := atomic.AddInt32(&errCount, 1)
			if ec > 10 {
				panic(err)
			}
		}
	}
}

func queueFor(prefix string) string {
	hasher.Reset()
	hasher.Write([]byte(prefix))
	hash := hasher.Sum32()
	return fmt.Sprintf("%s%d", *fQueue, int(hash)%*fNumQueues)
}

func day(prefix string) {
	log.Println(prefix)
	q := storage.Query{
		Delimiter: "/",
		// TODO - validate.
		Prefix: prefix,
	}
	it := bucket.Objects(context.Background(), &q)
	var wg sync.WaitGroup
	wg.Add(1)
	go postDay(&wg, queueFor(prefix), it)
	log.Println("Waiting")
	wg.Wait()
}

func month(prefix string) {
	log.Println(prefix)
	q := storage.Query{
		Delimiter: "/",
		// TODO - validate.
		Prefix: prefix,
	}
	it := bucket.Objects(context.Background(), &q)

	var wg sync.WaitGroup
	for o, err := it.Next(); err != iterator.Done; o, err = it.Next() {
		if err != nil {
			log.Println(err)
		}
		//		log.Printf("%+v\n", o)
		if o.Prefix != "" {
			q := storage.Query{
				Delimiter: "/",
				// TODO - validate.
				Prefix: o.Prefix,
			}
			it := bucket.Objects(context.Background(), &q)
			queue := queueFor(o.Prefix)
			wg.Add(1)
			go postDay(&wg, queue, it)
		} else {
			log.Println("Skipping: ", o.Name)
		}
	}
	log.Println("Waiting")
	wg.Wait()
}

func main() {
	flag.Parse()

	var err error
	storageClient, err = storage.NewClient(context.Background())
	if err != nil {
		log.Println(err)
		panic(err)
	}

	bucket = storageClient.Bucket(*fBucket)
	attr, err := bucket.Attrs(context.Background())
	if err != nil {
		log.Println(err)
		panic(err)
	}
	log.Println(attr)

	if *fMonth != "" {
		month(*fExper + "/" + *fMonth + "/")
	} else if *fDay != "" {
		day(*fExper + "/" + *fDay + "/")
	}
}
<commit_msg>Small improvements in queue assignment and code reuse<commit_after>// Package main defines a command line tool for submitting date
// ranges for reprocessing
package main

// TODO - note about setting up batch table and giving permission.

/*
Strategies...
  1. Work from a prefix, or range of prefixes.
  2. Work from a date range
  3. Work from a month prefix, but explicitly iterate over days.
      maybe use a separate goroutine for each date?

Usage:

*/

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"google.golang.org/api/iterator"

	"cloud.google.com/go/storage"
)

var (
	fProject   = flag.String("project", "mlab-oti", "Project containing queues.")
	fQueue     = flag.String("queue", "etl-ndt-batch-", "Base of queue name.")
	fNumQueues = flag.Int("num_queues", 5, "Number of queues.  Normally determined by listing queues.")
	fBucket    = flag.String("bucket", "archive-mlab-oti", "Source bucket.")
	fExper     = flag.String("experiment", "ndt", "Experiment prefix, trailing slash optional")
	fMonth     = flag.String("month", "", "Single month spec, as YYYY/MM")
	fDay       = flag.String("day", "", "Single day spec, as YYYY/MM/DD")

	errCount      int32
	storageClient *storage.Client
	bucket        *storage.BucketHandle
)

func init() {
	// Always prepend the filename and line number.
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func postOne(queue string, bucket string, fn string) error {
	reqStr := fmt.Sprintf("http://queue-pusher-dot-%s.appspot.com/receiver?queue=%s&filename=gs://%s/%s", *fProject, queue, bucket, fn)
	resp, err := http.Get(reqStr)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return errors.New("http error: " + resp.Status)
	}

	return nil
}

// Post all items in an ObjectIterator into specific
// queue.
func postDay(wg *sync.WaitGroup, queue string, it *storage.ObjectIterator) {
	if wg != nil {
		defer wg.Done()
	}
	for o, err := it.Next(); err != iterator.Done; o, err = it.Next() {
		if err != nil {
			// TODO - should this retry?
			log.Println(err)
			ec := atomic.AddInt32(&errCount, 1)
			if ec > 10 {
				panic(err)
			}
		}

		err = postOne(queue, *fBucket, o.Name)
		if err != nil {
			// TODO - should this retry?
			log.Println(err)
			ec := atomic.AddInt32(&errCount, 1)
			if ec > 10 {
				panic(err)
			}
		}
	}
}

// Initially this used a hash, but using day ordinal is better
// as it distributes across the queues more evenly.
func queueFor(date time.Time) string {
	day := date.Unix() / (24 * 60 * 60)
	return fmt.Sprintf("%s%d", *fQueue, int(day)%*fNumQueues)
}

func dateFromPrefixOrDie(prefix string) time.Time {
	date, err := time.Parse("ndt/2006/01/02/", prefix)
	if err != nil {
		log.Println(err)
		panic(err)
	}
	return date
}

func day(wg *sync.WaitGroup, prefix string) {
	queue := queueFor(dateFromPrefixOrDie(prefix))
	log.Println("Adding ", prefix, " to ", queue)
	q := storage.Query{
		Delimiter: "/",
		// TODO - validate.
		Prefix: prefix,
	}
	it := bucket.Objects(context.Background(), &q)
	go postDay(wg, queue, it)
}

func month(prefix string) {
	q := storage.Query{
		Delimiter: "/",
		// TODO - validate.
		Prefix: prefix,
	}
	it := bucket.Objects(context.Background(), &q)

	var wg sync.WaitGroup
	for o, err := it.Next(); err != iterator.Done; o, err = it.Next() {
		if err != nil {
			log.Println(err)
		} else if o.Prefix != "" {
			wg.Add(1)
			day(&wg, o.Prefix)
		} else {
			log.Println("Skipping: ", o.Name)
		}
	}
	wg.Wait()
}

func main() {
	flag.Parse()

	var err error
	storageClient, err = storage.NewClient(context.Background())
	if err != nil {
		log.Println(err)
		panic(err)
	}

	bucket = storageClient.Bucket(*fBucket)
	_, err = bucket.Attrs(context.Background())
	if err != nil {
		log.Println(err)
		panic(err)
	}

	if *fMonth != "" {
		month(*fExper + "/" + *fMonth + "/")
	} else if *fDay != "" {
		day(nil, *fExper+"/"+*fDay+"/")
	}
}
<|endoftext|>


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

// root2npy converts the content of a ROOT TTree to a NumPy data file.
//
//  Usage of root2npy:
//   -f string
//     	path to input ROOT file name
//   -o string
//     	path to output npz file name (default "output.npz")
//   -t string
//     	name of the tree to convert (default "tree")
//
package main

import (
	"archive/zip"
	"bytes"
	"flag"
	"io"
	"log"
	"os"
	"reflect"

	"github.com/sbinet/npyio"

	"go-hep.org/x/hep/rootio"
)

func main() {
	log.SetPrefix("root2npy: ")
	log.SetFlags(0)

	fname := flag.String("f", "", "path to input ROOT file name")
	oname := flag.String("o", "output.npz", "path to output npz file name")
	tname := flag.String("t", "tree", "name of the tree to convert")

	flag.Parse()

	if *fname == "" {
		flag.Usage()
		log.Fatalf("missing input ROOT filename argument")
	}

	f, err := rootio.Open(*fname)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	obj, ok := f.Get(*tname)
	if !ok {
		log.Fatalf("no object named %q in file %q", *tname, *fname)
	}

	tree, ok := obj.(rootio.Tree)
	if !ok {
		log.Fatalf("object %q in file %q is not a rootio.Tree", *tname, *fname)
	}

	var nt = ntuple{n: tree.Entries()}
	log.Printf("scanning leaves...")
	for _, leaf := range tree.Leaves() {
		if leaf.Kind() == reflect.String {
			nt.add(leaf.Name(), leaf)
			continue
		}
		if leaf.Class() == "TLeafElement" { // FIXME(sbinet): find a better, type-safe way
			log.Printf(">>> %q %v not supported", leaf.Name(), leaf.Class())
			continue
		}
		if leaf.LeafCount() != nil {
			log.Printf(">>> %q []%v not supported", leaf.Name(), leaf.TypeName())
			continue
		}
		nt.add(leaf.Name(), leaf)
	}
	log.Printf("scanning leaves... [done]")

	sc, err := rootio.NewScannerVars(tree, nt.args...)
	if err != nil {
		log.Fatal(err)
	}
	defer sc.Close()

	for sc.Next() {
		err = sc.Scan(nt.vars...)
		if err != nil {
			log.Fatal(err)
		}
		nt.fill()
	}

	out, err := os.Create(*oname)
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()

	npz := zip.NewWriter(out)
	defer npz.Close()

	for _, col := range nt.cols {
		buf := new(bytes.Buffer)
		err = npyio.Write(buf, col.slice.Interface())
		if err != nil {
			log.Fatalf("error writing %q: %v\n", col.name, err)
		}

		wz, err := npz.Create(col.name)
		if err != nil {
			log.Fatalf("error creating %q: %v\n", col.name, err)
		}

		_, err = io.Copy(wz, buf)
		if err != nil {
			log.Fatal(err)
		}
	}

	err = npz.Flush()
	if err != nil {
		log.Fatal(err)
	}

	err = npz.Close()
	if err != nil {
		log.Fatal(err)
	}

	err = out.Close()
	if err != nil {
		log.Fatal(err)
	}
}

type ntuple struct {
	n    int64
	cols []column
	args []rootio.ScanVar
	vars []interface{}
}

func (nt *ntuple) add(name string, leaf rootio.Leaf) {
	n := len(nt.cols)
	nt.cols = append(nt.cols, newColumn(name, leaf, nt.n))
	col := &nt.cols[n]
	nt.args = append(nt.args, rootio.ScanVar{Name: name, Type: col.etype})
	nt.vars = append(nt.vars, col.data.Addr().Interface())
}

func (nt *ntuple) fill() {
	for i := range nt.cols {
		col := &nt.cols[i]
		col.fill()
	}
}

type column struct {
	name  string
	i     int64
	leaf  rootio.Leaf
	etype reflect.Type
	shape []int
	data  reflect.Value
	slice reflect.Value
}

func newColumn(name string, leaf rootio.Leaf, n int64) column {
	etype := leaf.Type()
	shape := []int{int(n)}
	if leaf.Len() > 1 && leaf.Kind() != reflect.String {
		etype = reflect.ArrayOf(leaf.Len(), etype)
		shape = append(shape, leaf.Len())
	}
	rtype := reflect.SliceOf(etype)
	return column{
		name:  name,
		i:     0,
		leaf:  leaf,
		etype: etype,
		shape: shape,
		data:  reflect.New(etype).Elem(),
		slice: reflect.MakeSlice(rtype, int(n), int(n)),
	}
}

func (col *column) fill() {
	col.slice.Index(int(col.i)).Set(col.data)
	col.i++
}
<commit_msg>cmd/root2npy: improve documentation<commit_after>// Copyright 2017 The go-hep Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// root2npy converts the content of a ROOT TTree to a NumPy data file.
//
//  Usage of root2npy:
//   -f string
//     	path to input ROOT file name
//   -o string
//     	path to output npz file name (default "output.npz")
//   -t string
//     	name of the tree to convert (default "tree")
//
// Example:
//
//  $> root2npy -f $GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root -t tree -o output.npz
//  $> python2 -c 'import sys, numpy as np; print(dict(np.load(sys.argv[1])))' ./output.npz
//  {'one':   array([1, 2, 3, 4], dtype=int32),
//   'two':   array([ 1.10000002,  2.20000005,  3.29999995,  4.4000001 ], dtype=float32),
//   'three': array([u'uno', u'dos', u'tres', u'quatro'], dtype='<U6')}
//
//  $> python3 -c 'import sys, numpy as np; print(dict(np.load(sys.argv[1])))' ./output.npz
//  {'one':   array([1, 2, 3, 4], dtype=int32),
//   'two':   array([ 1.10000002,  2.20000005,  3.29999995,  4.4000001 ], dtype=float32),
//   'three': array(['uno', 'dos', 'tres', 'quatro'], dtype='<U6')}
//
//  $> go get github.com/sbinet/npyio/cmd/npyio-ls
//  $> npyio-ls ./output.npz
//  ================================================================================
//  file: ./output.npz
//  entry: one
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[4]}}
//  data = [1 2 3 4]
//
//  entry: two
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<f4, Fortran:false, Shape:[4]}}
//  data = [1.1 2.2 3.3 4.4]
//
//  entry: three
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<U6, Fortran:false, Shape:[4]}}
//  data = [uno dos tres quatro]
//
//  $> root-ls -t $GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root
//  === [$GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root] ===
//  version: 60600
//  TTree   tree      fake data (entries=4)
//    one   "one/I"   TBranch
//    two   "two/F"   TBranch
//    three "three/C" TBranch
//
// If you have a 10-events tree with a branch "doubles" containing an array of 3 float64,
// root2npy will convert it to a NumPy data file containing a NumPy array with a shape (10,3).
//
// Example:
//
//  $> root-ls -t $GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root
//  === [$GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root] ===
//  version: 60806
//  TTree          tree                 my tree title (entries=100)
//    Int32        "Int32/I"            TBranch
//    Int64        "Int64/L"            TBranch
//    UInt32       "UInt32/i"           TBranch
//    UInt64       "UInt64/l"           TBranch
//    Float32      "Float32/F"          TBranch
//    Float64      "Float64/D"          TBranch
//    Str          "Str/C"              TBranch
//    ArrayInt32   "ArrayInt32[10]/I"   TBranch
//    ArrayInt64   "ArrayInt64[10]/L"   TBranch
//    ArrayUInt32  "ArrayInt32[10]/i"   TBranch
//    ArrayUInt64  "ArrayInt64[10]/l"   TBranch
//    ArrayFloat32 "ArrayFloat32[10]/F" TBranch
//    ArrayFloat64 "ArrayFloat64[10]/D" TBranch
//    N            "N/I"                TBranch
//    SliceInt32   "SliceInt32[N]/I"    TBranch
//    SliceInt64   "SliceInt64[N]/L"    TBranch
//    SliceUInt32  "SliceInt32[N]/i"    TBranch
//    SliceUInt64  "SliceInt64[N]/l"    TBranch
//    SliceFloat32 "SliceFloat32[N]/F"  TBranch
//    SliceFloat64 "SliceFloat64[N]/D"  TBranch
//
//  $> root2npy $GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root
//  root2npy: scanning leaves...
//  root2npy: >>> "SliceInt32" []int32 not supported
//  root2npy: >>> "SliceInt64" []int64 not supported
//  root2npy: >>> "SliceInt32" []int32 not supported
//  root2npy: >>> "SliceInt64" []int64 not supported
//  root2npy: >>> "SliceFloat32" []float32 not supported
//  root2npy: >>> "SliceFloat64" []float64 not supported
//  root2npy: scanning leaves... [done]
//
//  $> npyio-ls ./output.npz
//  ================================================================================
//  file: ./output.npz
//  entry: Int32
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  entry: Int64
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i8, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  [...]
//
//  entry: Float64
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<f8, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  entry: Str
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<U7, Fortran:false, Shape:[100]}}
//  data = [evt-000 evt-001 evt-002 evt-003 evt-004 evt-005 evt-006 evt-007 ...
//  evt-092 evt-093 evt-094 evt-095 evt-096 evt-097 evt-098 evt-099]
//
//  entry: ArrayInt32
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[100 10]}}
//  data = [0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2 ...
//  ... 97 98 98 98 98 98 98 98 98 98 98 99 99 99 99 99 99 99 99 99 99]
//
//  [...]
//
package main

import (
	"archive/zip"
	"bytes"
	"flag"
	"io"
	"log"
	"os"
	"reflect"

	"github.com/sbinet/npyio"

	"go-hep.org/x/hep/rootio"
)

func main() {
	log.SetPrefix("root2npy: ")
	log.SetFlags(0)

	fname := flag.String("f", "", "path to input ROOT file name")
	oname := flag.String("o", "output.npz", "path to output npz file name")
	tname := flag.String("t", "tree", "name of the tree to convert")

	flag.Parse()

	if *fname == "" {
		flag.Usage()
		log.Fatalf("missing input ROOT filename argument")
	}

	f, err := rootio.Open(*fname)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	obj, ok := f.Get(*tname)
	if !ok {
		log.Fatalf("no object named %q in file %q", *tname, *fname)
	}

	tree, ok := obj.(rootio.Tree)
	if !ok {
		log.Fatalf("object %q in file %q is not a rootio.Tree", *tname, *fname)
	}

	var nt = ntuple{n: tree.Entries()}
	log.Printf("scanning leaves...")
	for _, leaf := range tree.Leaves() {
		if leaf.Kind() == reflect.String {
			nt.add(leaf.Name(), leaf)
			continue
		}
		if leaf.Class() == "TLeafElement" { // FIXME(sbinet): find a better, type-safe way
			log.Printf(">>> %q %v not supported", leaf.Name(), leaf.Class())
			continue
		}
		if leaf.LeafCount() != nil {
			log.Printf(">>> %q []%v not supported", leaf.Name(), leaf.TypeName())
			continue
		}
		nt.add(leaf.Name(), leaf)
	}
	log.Printf("scanning leaves... [done]")

	sc, err := rootio.NewScannerVars(tree, nt.args...)
	if err != nil {
		log.Fatal(err)
	}
	defer sc.Close()

	for sc.Next() {
		err = sc.Scan(nt.vars...)
		if err != nil {
			log.Fatal(err)
		}
		nt.fill()
	}

	out, err := os.Create(*oname)
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()

	npz := zip.NewWriter(out)
	defer npz.Close()

	for _, col := range nt.cols {
		buf := new(bytes.Buffer)
		err = npyio.Write(buf, col.slice.Interface())
		if err != nil {
			log.Fatalf("error writing %q: %v\n", col.name, err)
		}

		wz, err := npz.Create(col.name)
		if err != nil {
			log.Fatalf("error creating %q: %v\n", col.name, err)
		}

		_, err = io.Copy(wz, buf)
		if err != nil {
			log.Fatal(err)
		}
	}

	err = npz.Flush()
	if err != nil {
		log.Fatal(err)
	}

	err = npz.Close()
	if err != nil {
		log.Fatal(err)
	}

	err = out.Close()
	if err != nil {
		log.Fatal(err)
	}
}

type ntuple struct {
	n    int64
	cols []column
	args []rootio.ScanVar
	vars []interface{}
}

func (nt *ntuple) add(name string, leaf rootio.Leaf) {
	n := len(nt.cols)
	nt.cols = append(nt.cols, newColumn(name, leaf, nt.n))
	col := &nt.cols[n]
	nt.args = append(nt.args, rootio.ScanVar{Name: name, Type: col.etype})
	nt.vars = append(nt.vars, col.data.Addr().Interface())
}

func (nt *ntuple) fill() {
	for i := range nt.cols {
		col := &nt.cols[i]
		col.fill()
	}
}

type column struct {
	name  string
	i     int64
	leaf  rootio.Leaf
	etype reflect.Type
	shape []int
	data  reflect.Value
	slice reflect.Value
}

func newColumn(name string, leaf rootio.Leaf, n int64) column {
	etype := leaf.Type()
	shape := []int{int(n)}
	if leaf.Len() > 1 && leaf.Kind() != reflect.String {
		etype = reflect.ArrayOf(leaf.Len(), etype)
		shape = append(shape, leaf.Len())
	}
	rtype := reflect.SliceOf(etype)
	return column{
		name:  name,
		i:     0,
		leaf:  leaf,
		etype: etype,
		shape: shape,
		data:  reflect.New(etype).Elem(),
		slice: reflect.MakeSlice(rtype, int(n), int(n)),
	}
}

func (col *column) fill() {
	col.slice.Index(int(col.i)).Set(col.data)
	col.i++
}
<|endoftext|>


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

// root2npy converts the content of a ROOT TTree to a NumPy data file.
//
//  Usage of root2npy:
//   -f string
//     	path to input ROOT file name
//   -o string
//     	path to output npz file name (default "output.npz")
//   -t string
//     	name of the tree to convert (default "tree")
//
// Example:
//
//  $> root2npy -f $GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root -t tree -o output.npz
//  $> python2 -c 'import sys, numpy as np; print(dict(np.load(sys.argv[1])))' ./output.npz
//  {'one':   array([1, 2, 3, 4], dtype=int32),
//   'two':   array([ 1.10000002,  2.20000005,  3.29999995,  4.4000001 ], dtype=float32),
//   'three': array([u'uno', u'dos', u'tres', u'quatro'], dtype='<U6')}
//
//  $> python3 -c 'import sys, numpy as np; print(dict(np.load(sys.argv[1])))' ./output.npz
//  {'one':   array([1, 2, 3, 4], dtype=int32),
//   'two':   array([ 1.10000002,  2.20000005,  3.29999995,  4.4000001 ], dtype=float32),
//   'three': array(['uno', 'dos', 'tres', 'quatro'], dtype='<U6')}
//
//  $> go get github.com/sbinet/npyio/cmd/npyio-ls
//  $> npyio-ls ./output.npz
//  ================================================================================
//  file: ./output.npz
//  entry: one
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[4]}}
//  data = [1 2 3 4]
//
//  entry: two
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<f4, Fortran:false, Shape:[4]}}
//  data = [1.1 2.2 3.3 4.4]
//
//  entry: three
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<U6, Fortran:false, Shape:[4]}}
//  data = [uno dos tres quatro]
//
//  $> root-ls -t $GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root
//  === [$GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root] ===
//  version: 60600
//  TTree   tree      fake data (entries=4)
//    one   "one/I"   TBranch
//    two   "two/F"   TBranch
//    three "three/C" TBranch
//
// If you have a 10-events tree with a branch "doubles" containing an array of 3 float64,
// root2npy will convert it to a NumPy data file containing a NumPy array with a shape (10,3).
//
// Example:
//
//  $> root-ls -t $GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root
//  === [$GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root] ===
//  version: 60806
//  TTree          tree                 my tree title (entries=100)
//    Int32        "Int32/I"            TBranch
//    Int64        "Int64/L"            TBranch
//    UInt32       "UInt32/i"           TBranch
//    UInt64       "UInt64/l"           TBranch
//    Float32      "Float32/F"          TBranch
//    Float64      "Float64/D"          TBranch
//    Str          "Str/C"              TBranch
//    ArrayInt32   "ArrayInt32[10]/I"   TBranch
//    ArrayInt64   "ArrayInt64[10]/L"   TBranch
//    ArrayUInt32  "ArrayInt32[10]/i"   TBranch
//    ArrayUInt64  "ArrayInt64[10]/l"   TBranch
//    ArrayFloat32 "ArrayFloat32[10]/F" TBranch
//    ArrayFloat64 "ArrayFloat64[10]/D" TBranch
//    N            "N/I"                TBranch
//    SliceInt32   "SliceInt32[N]/I"    TBranch
//    SliceInt64   "SliceInt64[N]/L"    TBranch
//    SliceUInt32  "SliceInt32[N]/i"    TBranch
//    SliceUInt64  "SliceInt64[N]/l"    TBranch
//    SliceFloat32 "SliceFloat32[N]/F"  TBranch
//    SliceFloat64 "SliceFloat64[N]/D"  TBranch
//
//  $> root2npy $GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root
//  root2npy: scanning leaves...
//  root2npy: >>> "SliceInt32" []int32 not supported
//  root2npy: >>> "SliceInt64" []int64 not supported
//  root2npy: >>> "SliceInt32" []int32 not supported
//  root2npy: >>> "SliceInt64" []int64 not supported
//  root2npy: >>> "SliceFloat32" []float32 not supported
//  root2npy: >>> "SliceFloat64" []float64 not supported
//  root2npy: scanning leaves... [done]
//
//  $> npyio-ls ./output.npz
//  ================================================================================
//  file: ./output.npz
//  entry: Int32
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  entry: Int64
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i8, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  [...]
//
//  entry: Float64
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<f8, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  entry: Str
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<U7, Fortran:false, Shape:[100]}}
//  data = [evt-000 evt-001 evt-002 evt-003 evt-004 evt-005 evt-006 evt-007 ...
//  evt-092 evt-093 evt-094 evt-095 evt-096 evt-097 evt-098 evt-099]
//
//  entry: ArrayInt32
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[100 10]}}
//  data = [0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2 ...
//  ... 97 98 98 98 98 98 98 98 98 98 98 99 99 99 99 99 99 99 99 99 99]
//
//  [...]
//
package main

import (
	"archive/zip"
	"bytes"
	"flag"
	"io"
	"log"
	"os"
	"reflect"

	"github.com/sbinet/npyio"

	"go-hep.org/x/hep/rootio"
)

func main() {
	log.SetPrefix("root2npy: ")
	log.SetFlags(0)

	fname := flag.String("f", "", "path to input ROOT file name")
	oname := flag.String("o", "output.npz", "path to output npz file name")
	tname := flag.String("t", "tree", "name of the tree to convert")

	flag.Parse()

	if *fname == "" {
		flag.Usage()
		log.Fatalf("missing input ROOT filename argument")
	}

	f, err := rootio.Open(*fname)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	obj, ok := f.Get(*tname)
	if !ok {
		log.Fatalf("no object named %q in file %q", *tname, *fname)
	}

	tree, ok := obj.(rootio.Tree)
	if !ok {
		log.Fatalf("object %q in file %q is not a rootio.Tree", *tname, *fname)
	}

	var nt = ntuple{n: tree.Entries()}
	log.Printf("scanning leaves...")
	for _, leaf := range tree.Leaves() {
		if leaf.Kind() == reflect.String {
			nt.add(leaf.Name(), leaf)
			continue
		}
		if leaf.Class() == "TLeafElement" { // FIXME(sbinet): find a better, type-safe way
			log.Printf(">>> %q %v not supported", leaf.Name(), leaf.Class())
			continue
		}
		if leaf.LeafCount() != nil {
			log.Printf(">>> %q []%v not supported", leaf.Name(), leaf.TypeName())
			continue
		}
		nt.add(leaf.Name(), leaf)
	}
	log.Printf("scanning leaves... [done]")

	sc, err := rootio.NewScannerVars(tree, nt.args...)
	if err != nil {
		log.Fatal(err)
	}
	defer sc.Close()

	for sc.Next() {
		err = sc.Scan(nt.vars...)
		if err != nil {
			log.Fatal(err)
		}
		nt.fill()
	}

	out, err := os.Create(*oname)
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()

	npz := zip.NewWriter(out)
	defer npz.Close()

	for _, col := range nt.cols {
		buf := new(bytes.Buffer)
		err = npyio.Write(buf, col.slice.Interface())
		if err != nil {
			log.Fatalf("error writing %q: %v\n", col.name, err)
		}

		wz, err := npz.Create(col.name)
		if err != nil {
			log.Fatalf("error creating %q: %v\n", col.name, err)
		}

		_, err = io.Copy(wz, buf)
		if err != nil {
			log.Fatal(err)
		}
	}

	err = npz.Flush()
	if err != nil {
		log.Fatal(err)
	}

	err = npz.Close()
	if err != nil {
		log.Fatal(err)
	}

	err = out.Close()
	if err != nil {
		log.Fatal(err)
	}
}

type ntuple struct {
	n    int64
	cols []column
	args []rootio.ScanVar
	vars []interface{}
}

func (nt *ntuple) add(name string, leaf rootio.Leaf) {
	n := len(nt.cols)
	nt.cols = append(nt.cols, newColumn(name, leaf, nt.n))
	col := &nt.cols[n]
	nt.args = append(nt.args, rootio.ScanVar{Name: name, Type: col.etype})
	nt.vars = append(nt.vars, col.data.Addr().Interface())
}

func (nt *ntuple) fill() {
	for i := range nt.cols {
		col := &nt.cols[i]
		col.fill()
	}
}

type column struct {
	name  string
	i     int64
	leaf  rootio.Leaf
	etype reflect.Type
	shape []int
	data  reflect.Value
	slice reflect.Value
}

func newColumn(name string, leaf rootio.Leaf, n int64) column {
	etype := leaf.Type()
	shape := []int{int(n)}
	if leaf.Len() > 1 && leaf.Kind() != reflect.String {
		etype = reflect.ArrayOf(leaf.Len(), etype)
		shape = append(shape, leaf.Len())
	}
	rtype := reflect.SliceOf(etype)
	return column{
		name:  name,
		i:     0,
		leaf:  leaf,
		etype: etype,
		shape: shape,
		data:  reflect.New(etype).Elem(),
		slice: reflect.MakeSlice(rtype, int(n), int(n)),
	}
}

func (col *column) fill() {
	col.slice.Index(int(col.i)).Set(col.data)
	col.i++
}
<commit_msg>cmd/root2npy: add link to NumPy data file format<commit_after>// Copyright 2017 The go-hep Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// root2npy converts the content of a ROOT TTree to a NumPy data file.
//
//  Usage of root2npy:
//   -f string
//     	path to input ROOT file name
//   -o string
//     	path to output npz file name (default "output.npz")
//   -t string
//     	name of the tree to convert (default "tree")
//
// The NumPy data file format is described here:
//
//  https://docs.scipy.org/doc/numpy/neps/npy-format.html
//
// Example:
//
//  $> root2npy -f $GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root -t tree -o output.npz
//  $> python2 -c 'import sys, numpy as np; print(dict(np.load(sys.argv[1])))' ./output.npz
//  {'one':   array([1, 2, 3, 4], dtype=int32),
//   'two':   array([ 1.10000002,  2.20000005,  3.29999995,  4.4000001 ], dtype=float32),
//   'three': array([u'uno', u'dos', u'tres', u'quatro'], dtype='<U6')}
//
//  $> python3 -c 'import sys, numpy as np; print(dict(np.load(sys.argv[1])))' ./output.npz
//  {'one':   array([1, 2, 3, 4], dtype=int32),
//   'two':   array([ 1.10000002,  2.20000005,  3.29999995,  4.4000001 ], dtype=float32),
//   'three': array(['uno', 'dos', 'tres', 'quatro'], dtype='<U6')}
//
//  $> go get github.com/sbinet/npyio/cmd/npyio-ls
//  $> npyio-ls ./output.npz
//  ================================================================================
//  file: ./output.npz
//  entry: one
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[4]}}
//  data = [1 2 3 4]
//
//  entry: two
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<f4, Fortran:false, Shape:[4]}}
//  data = [1.1 2.2 3.3 4.4]
//
//  entry: three
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<U6, Fortran:false, Shape:[4]}}
//  data = [uno dos tres quatro]
//
//  $> root-ls -t $GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root
//  === [$GOPATH/src/go-hep.org/x/hep/rootio/testdata/simple.root] ===
//  version: 60600
//  TTree   tree      fake data (entries=4)
//    one   "one/I"   TBranch
//    two   "two/F"   TBranch
//    three "three/C" TBranch
//
// If you have a 10-events tree with a branch "doubles" containing an array of 3 float64,
// root2npy will convert it to a NumPy data file containing a NumPy array with a shape (10,3).
//
// Example:
//
//  $> root-ls -t $GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root
//  === [$GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root] ===
//  version: 60806
//  TTree          tree                 my tree title (entries=100)
//    Int32        "Int32/I"            TBranch
//    Int64        "Int64/L"            TBranch
//    UInt32       "UInt32/i"           TBranch
//    UInt64       "UInt64/l"           TBranch
//    Float32      "Float32/F"          TBranch
//    Float64      "Float64/D"          TBranch
//    Str          "Str/C"              TBranch
//    ArrayInt32   "ArrayInt32[10]/I"   TBranch
//    ArrayInt64   "ArrayInt64[10]/L"   TBranch
//    ArrayUInt32  "ArrayInt32[10]/i"   TBranch
//    ArrayUInt64  "ArrayInt64[10]/l"   TBranch
//    ArrayFloat32 "ArrayFloat32[10]/F" TBranch
//    ArrayFloat64 "ArrayFloat64[10]/D" TBranch
//    N            "N/I"                TBranch
//    SliceInt32   "SliceInt32[N]/I"    TBranch
//    SliceInt64   "SliceInt64[N]/L"    TBranch
//    SliceUInt32  "SliceInt32[N]/i"    TBranch
//    SliceUInt64  "SliceInt64[N]/l"    TBranch
//    SliceFloat32 "SliceFloat32[N]/F"  TBranch
//    SliceFloat64 "SliceFloat64[N]/D"  TBranch
//
//  $> root2npy $GOPATH/src/go-hep.org/x/hep/rootio/testdata/small-flat-tree.root
//  root2npy: scanning leaves...
//  root2npy: >>> "SliceInt32" []int32 not supported
//  root2npy: >>> "SliceInt64" []int64 not supported
//  root2npy: >>> "SliceInt32" []int32 not supported
//  root2npy: >>> "SliceInt64" []int64 not supported
//  root2npy: >>> "SliceFloat32" []float32 not supported
//  root2npy: >>> "SliceFloat64" []float64 not supported
//  root2npy: scanning leaves... [done]
//
//  $> npyio-ls ./output.npz
//  ================================================================================
//  file: ./output.npz
//  entry: Int32
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  entry: Int64
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i8, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  [...]
//
//  entry: Float64
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<f8, Fortran:false, Shape:[100]}}
//  data = [0 1 2 3 4 5 6 7 8 9 10 11 ... 90 91 92 93 94 95 96 97 98 99]
//
//  entry: Str
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<U7, Fortran:false, Shape:[100]}}
//  data = [evt-000 evt-001 evt-002 evt-003 evt-004 evt-005 evt-006 evt-007 ...
//  evt-092 evt-093 evt-094 evt-095 evt-096 evt-097 evt-098 evt-099]
//
//  entry: ArrayInt32
//  npy-header: Header{Major:2, Minor:0, Descr:{Type:<i4, Fortran:false, Shape:[100 10]}}
//  data = [0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2 ...
//  ... 97 98 98 98 98 98 98 98 98 98 98 99 99 99 99 99 99 99 99 99 99]
//
//  [...]
//
package main

import (
	"archive/zip"
	"bytes"
	"flag"
	"io"
	"log"
	"os"
	"reflect"

	"github.com/sbinet/npyio"

	"go-hep.org/x/hep/rootio"
)

func main() {
	log.SetPrefix("root2npy: ")
	log.SetFlags(0)

	fname := flag.String("f", "", "path to input ROOT file name")
	oname := flag.String("o", "output.npz", "path to output npz file name")
	tname := flag.String("t", "tree", "name of the tree to convert")

	flag.Parse()

	if *fname == "" {
		flag.Usage()
		log.Fatalf("missing input ROOT filename argument")
	}

	f, err := rootio.Open(*fname)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	obj, ok := f.Get(*tname)
	if !ok {
		log.Fatalf("no object named %q in file %q", *tname, *fname)
	}

	tree, ok := obj.(rootio.Tree)
	if !ok {
		log.Fatalf("object %q in file %q is not a rootio.Tree", *tname, *fname)
	}

	var nt = ntuple{n: tree.Entries()}
	log.Printf("scanning leaves...")
	for _, leaf := range tree.Leaves() {
		if leaf.Kind() == reflect.String {
			nt.add(leaf.Name(), leaf)
			continue
		}
		if leaf.Class() == "TLeafElement" { // FIXME(sbinet): find a better, type-safe way
			log.Printf(">>> %q %v not supported", leaf.Name(), leaf.Class())
			continue
		}
		if leaf.LeafCount() != nil {
			log.Printf(">>> %q []%v not supported", leaf.Name(), leaf.TypeName())
			continue
		}
		nt.add(leaf.Name(), leaf)
	}
	log.Printf("scanning leaves... [done]")

	sc, err := rootio.NewScannerVars(tree, nt.args...)
	if err != nil {
		log.Fatal(err)
	}
	defer sc.Close()

	for sc.Next() {
		err = sc.Scan(nt.vars...)
		if err != nil {
			log.Fatal(err)
		}
		nt.fill()
	}

	out, err := os.Create(*oname)
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()

	npz := zip.NewWriter(out)
	defer npz.Close()

	for _, col := range nt.cols {
		buf := new(bytes.Buffer)
		err = npyio.Write(buf, col.slice.Interface())
		if err != nil {
			log.Fatalf("error writing %q: %v\n", col.name, err)
		}

		wz, err := npz.Create(col.name)
		if err != nil {
			log.Fatalf("error creating %q: %v\n", col.name, err)
		}

		_, err = io.Copy(wz, buf)
		if err != nil {
			log.Fatal(err)
		}
	}

	err = npz.Flush()
	if err != nil {
		log.Fatal(err)
	}

	err = npz.Close()
	if err != nil {
		log.Fatal(err)
	}

	err = out.Close()
	if err != nil {
		log.Fatal(err)
	}
}

type ntuple struct {
	n    int64
	cols []column
	args []rootio.ScanVar
	vars []interface{}
}

func (nt *ntuple) add(name string, leaf rootio.Leaf) {
	n := len(nt.cols)
	nt.cols = append(nt.cols, newColumn(name, leaf, nt.n))
	col := &nt.cols[n]
	nt.args = append(nt.args, rootio.ScanVar{Name: name, Type: col.etype})
	nt.vars = append(nt.vars, col.data.Addr().Interface())
}

func (nt *ntuple) fill() {
	for i := range nt.cols {
		col := &nt.cols[i]
		col.fill()
	}
}

type column struct {
	name  string
	i     int64
	leaf  rootio.Leaf
	etype reflect.Type
	shape []int
	data  reflect.Value
	slice reflect.Value
}

func newColumn(name string, leaf rootio.Leaf, n int64) column {
	etype := leaf.Type()
	shape := []int{int(n)}
	if leaf.Len() > 1 && leaf.Kind() != reflect.String {
		etype = reflect.ArrayOf(leaf.Len(), etype)
		shape = append(shape, leaf.Len())
	}
	rtype := reflect.SliceOf(etype)
	return column{
		name:  name,
		i:     0,
		leaf:  leaf,
		etype: etype,
		shape: shape,
		data:  reflect.New(etype).Elem(),
		slice: reflect.MakeSlice(rtype, int(n), int(n)),
	}
}

func (col *column) fill() {
	col.slice.Index(int(col.i)).Set(col.data)
	col.i++
}
<|endoftext|>


<commit_before>package term

// These are used for drawing "ASCII-art" boxes
const (
	TLCHAR = '┌' // top left
	TRCHAR = '┐' // top right
	BLCHAR = '└' // bottom left
	BRCHAR = '┘' // bottom right
	VCHAR  = '│' // vertical line, left side
	VCHAR2 = '│' // vertical line, right side
	HCHAR  = '─' // horizontal line
	HCHAR2 = '─' // horizontal bottom line
)

// Rect is a position, width and height
type Rect struct {
	X int
	Y int
	W int
	H int
}

// Box has one outer and one innner rectangle.
// This is useful when having margins that surrounds content.
type Box struct {
	frame *Rect // The rectangle around the box, for placement
	inner *Rect // The rectangle inside the box, for content
}

// Create a new Box / container.
func NewBox() *Box {
	return &Box{&Rect{0, 0, 0, 0}, &Rect{0, 0, 0, 0}}
}

// Place a Box at the center of the given container.
func (b *Box) Center(container *Box) {
	widthleftover := container.inner.W - b.frame.W
	heightleftover := container.inner.H - b.frame.H
	b.frame.X = container.inner.X + widthleftover/2
	b.frame.Y = container.inner.Y + heightleftover/2
}

// Place a Box so that it fills the entire given container.
func (b *Box) Fill(container *Box) {
	b.frame.X = container.inner.X
	b.frame.Y = container.inner.Y
	b.frame.W = container.inner.W
	b.frame.H = container.inner.H
}

// Place a Box inside a given container, with the given margins.
// Margins are given in number of characters.
func (b *Box) FillWithMargins(container *Box, margins int) {
	b.Fill(container)
	b.frame.X += margins
	b.frame.Y += margins
	b.frame.W -= margins * 2
	b.frame.H -= margins * 2
}

// Place a Box inside a given container, using the given percentage wise ratios.
// horizmarginp can for example be 0.1 for a 10% horizontal margin around
// the inner box. vertmarginp works similarly, but for the vertical margins.
func (b *Box) FillWithPercentageMargins(container *Box, horizmarginp float32, vertmarginp float32) {
	horizmargin := int(float32(container.inner.W) * horizmarginp)
	vertmargin := int(float32(container.inner.H) * vertmarginp)
	b.Fill(container)
	b.frame.X += horizmargin
	b.frame.Y += vertmargin
	b.frame.W -= horizmargin * 2
	b.frame.H -= vertmargin * 2
}

// Retrieves the position of the inner rectangle.
func (b *Box) GetContentPos() (int, int) {
	return b.inner.X, b.inner.Y
}

// Set the size of the Box to 1/3 of the size of the inner rectangle
// of the given container.
func (b *Box) SetThirdSize(container *Box) {
	b.frame.W = container.inner.W / 3
	b.frame.H = container.inner.H / 3
}

// Set the position of the Box to 1/3 of the size of the inner rectangle
// of the given container.
func (b *Box) SetThirdPlace(container *Box) {
	b.frame.X = container.inner.X + container.inner.W/3
	b.frame.Y = container.inner.Y + container.inner.H/3
}

// Place a Box so that it either fills the given
// container, or is placed 1/3 from the upper left edge,
// depending on how much space is left.
func (b *Box) SetNicePlacement(container *Box) {
	b.frame.X = container.inner.X
	b.frame.Y = container.inner.Y
	leftoverwidth := container.inner.W - b.frame.W
	leftoverheight := container.inner.H - b.frame.H
	if leftoverwidth > b.frame.W {
		b.frame.X += leftoverwidth / 3
	}
	if leftoverheight > b.frame.H {
		b.frame.Y += leftoverheight / 3
	}
}

// Place a Box within the given container.
func (b *Box) Place(container *Box) {
	b.frame.X = container.inner.X
	b.frame.Y = container.inner.Y
}

// Get the inner rectangle (content size + pos)
func (b *Box) GetInner() *Rect {
	return b.inner
}

// Get the outer frame (box size + pos)
func (b *Box) GetFrame() *Rect {
	return b.frame
}

// Set the inner rectangle (content size + pos)
func (b *Box) SetInner(r *Rect) {
	b.inner = r
}

// Set the outer frame (box size + pos)
func (b *Box) SetFrame(r *Rect) {
	b.frame = r
}
<commit_msg>Rounded boxes<commit_after>package term

// These are used for drawing "ASCII-art" boxes
const (
	TLCHAR = '╭' // top left
	TRCHAR = '╮' // top right
	BLCHAR = '╰' // bottom left
	BRCHAR = '╯' // bottom right
	VCHAR  = '│' // vertical line, left side
	VCHAR2 = '│' // vertical line, right side
	HCHAR  = '─' // horizontal line
	HCHAR2 = '─' // horizontal bottom line
)

// Rect is a position, width and height
type Rect struct {
	X int
	Y int
	W int
	H int
}

// Box has one outer and one innner rectangle.
// This is useful when having margins that surrounds content.
type Box struct {
	frame *Rect // The rectangle around the box, for placement
	inner *Rect // The rectangle inside the box, for content
}

// Create a new Box / container.
func NewBox() *Box {
	return &Box{&Rect{0, 0, 0, 0}, &Rect{0, 0, 0, 0}}
}

// Place a Box at the center of the given container.
func (b *Box) Center(container *Box) {
	widthleftover := container.inner.W - b.frame.W
	heightleftover := container.inner.H - b.frame.H
	b.frame.X = container.inner.X + widthleftover/2
	b.frame.Y = container.inner.Y + heightleftover/2
}

// Place a Box so that it fills the entire given container.
func (b *Box) Fill(container *Box) {
	b.frame.X = container.inner.X
	b.frame.Y = container.inner.Y
	b.frame.W = container.inner.W
	b.frame.H = container.inner.H
}

// Place a Box inside a given container, with the given margins.
// Margins are given in number of characters.
func (b *Box) FillWithMargins(container *Box, margins int) {
	b.Fill(container)
	b.frame.X += margins
	b.frame.Y += margins
	b.frame.W -= margins * 2
	b.frame.H -= margins * 2
}

// Place a Box inside a given container, using the given percentage wise ratios.
// horizmarginp can for example be 0.1 for a 10% horizontal margin around
// the inner box. vertmarginp works similarly, but for the vertical margins.
func (b *Box) FillWithPercentageMargins(container *Box, horizmarginp float32, vertmarginp float32) {
	horizmargin := int(float32(container.inner.W) * horizmarginp)
	vertmargin := int(float32(container.inner.H) * vertmarginp)
	b.Fill(container)
	b.frame.X += horizmargin
	b.frame.Y += vertmargin
	b.frame.W -= horizmargin * 2
	b.frame.H -= vertmargin * 2
}

// Retrieves the position of the inner rectangle.
func (b *Box) GetContentPos() (int, int) {
	return b.inner.X, b.inner.Y
}

// Set the size of the Box to 1/3 of the size of the inner rectangle
// of the given container.
func (b *Box) SetThirdSize(container *Box) {
	b.frame.W = container.inner.W / 3
	b.frame.H = container.inner.H / 3
}

// Set the position of the Box to 1/3 of the size of the inner rectangle
// of the given container.
func (b *Box) SetThirdPlace(container *Box) {
	b.frame.X = container.inner.X + container.inner.W/3
	b.frame.Y = container.inner.Y + container.inner.H/3
}

// Place a Box so that it either fills the given
// container, or is placed 1/3 from the upper left edge,
// depending on how much space is left.
func (b *Box) SetNicePlacement(container *Box) {
	b.frame.X = container.inner.X
	b.frame.Y = container.inner.Y
	leftoverwidth := container.inner.W - b.frame.W
	leftoverheight := container.inner.H - b.frame.H
	if leftoverwidth > b.frame.W {
		b.frame.X += leftoverwidth / 3
	}
	if leftoverheight > b.frame.H {
		b.frame.Y += leftoverheight / 3
	}
}

// Place a Box within the given container.
func (b *Box) Place(container *Box) {
	b.frame.X = container.inner.X
	b.frame.Y = container.inner.Y
}

// Get the inner rectangle (content size + pos)
func (b *Box) GetInner() *Rect {
	return b.inner
}

// Get the outer frame (box size + pos)
func (b *Box) GetFrame() *Rect {
	return b.frame
}

// Set the inner rectangle (content size + pos)
func (b *Box) SetInner(r *Rect) {
	b.inner = r
}

// Set the outer frame (box size + pos)
func (b *Box) SetFrame(r *Rect) {
	b.frame = r
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/moov-io/ach"
	"log"
	"os"
	"path/filepath"
	"runtime/pprof"
	"time"
)

var (
	fPath      = flag.String("fPath", "", "File Path")
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
)

// main creates an ACH File with 4 batches of SEC Code PPD.
// Each batch contains an EntryAddendaCount of 2500.
func main() {
	flag.Parse()

	path := filepath.Join(*fPath, time.Now().UTC().Format("200601021504")+".ach")
	write(path)
}

func write(path string) {
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	f, err := os.Create(path)
	if err != nil {
		fmt.Printf("%T: %s", err, err)
	}

	// To create a file
	fh := ach.NewFileHeader()
	fh.ImmediateDestination = "231380104"
	fh.ImmediateOrigin = "121042882"
	fh.FileCreationDate = time.Now()
	fh.ImmediateDestinationName = "Citadel"
	fh.ImmediateOriginName = "Wells Fargo"
	file := ach.NewFile()
	file.SetHeader(fh)

	// Create 4 Batches of SEC Code PPD
	for i := 0; i < 4; i++ {
		bh := ach.NewBatchHeader()
		bh.ServiceClassCode = 200
		bh.CompanyName = "Wells Fargo"
		bh.CompanyIdentification = "121042882"
		bh.StandardEntryClassCode = "PPD"
		bh.CompanyEntryDescription = "Trans. Description"
		bh.EffectiveEntryDate = time.Now().AddDate(0, 0, 1)
		bh.ODFIIdentification = "121042882"

		batch, _ := ach.NewBatch(bh)

		// Create Entry
		entrySeq := 0
		for i := 0; i < 1250; i++ {
			entrySeq = entrySeq + 1

			entryEntrySeq := ach.NewEntryDetail()
			entryEntrySeq.TransactionCode = 22
			entryEntrySeq.SetRDFI("231380104")
			entryEntrySeq.DFIAccountNumber = "81967038518"
			entryEntrySeq.Amount = 100000
			//entryEntrySeq.IndividualName = randomdata.FullName(randomdata.RandomGender)
			entryEntrySeq.IndividualName = "Steven Tander"
			entryEntrySeq.SetTraceNumber(bh.ODFIIdentification, entrySeq)
			//entryEntrySeq.IdentificationNumber = "#" + randomdata.RandStringRunes(13) + "#"
			entryEntrySeq.IdentificationNumber = "#83738AB#"
			entryEntrySeq.Category = ach.CategoryForward

			// Add addenda record for an entry
			addendaEntrySeq := ach.NewAddenda05()
			addendaEntrySeq.PaymentRelatedInformation = "bonus pay for amazing work on #OSS"
			entryEntrySeq.AddAddenda(addendaEntrySeq)

			// Add entries
			batch.AddEntry(entryEntrySeq)

		}

		// Create the batch.
		if err := batch.Create(); err != nil {
			fmt.Printf("%T: %s", err, err)
		}

		// Add batch to the file
		file.AddBatch(batch)
	}

	// Create the file
	if err := file.Create(); err != nil {
		fmt.Printf("%T: %s", err, err)
	}

	// Write to a file
	w := ach.NewWriter(f)
	if err := w.Write(file); err != nil {
		fmt.Printf("%T: %s", err, err)
	}
	w.Flush()
	if err := f.Close(); err != nil {
		fmt.Println(err.Error())
	}
}
<commit_msg>cmd/writeACH: validate file before writing<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/moov-io/ach"
	"log"
	"os"
	"path/filepath"
	"runtime/pprof"
	"time"
)

var (
	fPath      = flag.String("fPath", "", "File Path")
	cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
)

// main creates an ACH File with 4 batches of SEC Code PPD.
// Each batch contains an EntryAddendaCount of 2500.
func main() {
	flag.Parse()

	path := filepath.Join(*fPath, time.Now().UTC().Format("200601021504")+".ach")
	write(path)
}

func write(path string) {
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	f, err := os.Create(path)
	if err != nil {
		fmt.Printf("%T: %s", err, err)
	}

	// To create a file
	fh := ach.NewFileHeader()
	fh.ImmediateDestination = "231380104"
	fh.ImmediateOrigin = "121042882"
	fh.FileCreationDate = time.Now()
	fh.ImmediateDestinationName = "Citadel"
	fh.ImmediateOriginName = "Wells Fargo"
	file := ach.NewFile()
	file.SetHeader(fh)

	// Create 4 Batches of SEC Code PPD
	for i := 0; i < 4; i++ {
		bh := ach.NewBatchHeader()
		bh.ServiceClassCode = 200
		bh.CompanyName = "Wells Fargo"
		bh.CompanyIdentification = "121042882"
		bh.StandardEntryClassCode = "PPD"
		bh.CompanyEntryDescription = "Trans. Description"
		bh.EffectiveEntryDate = time.Now().AddDate(0, 0, 1)
		bh.ODFIIdentification = "121042882"

		batch, _ := ach.NewBatch(bh)

		// Create Entry
		entrySeq := 0
		for i := 0; i < 1250; i++ {
			entrySeq = entrySeq + 1

			entryEntrySeq := ach.NewEntryDetail()
			entryEntrySeq.TransactionCode = 22
			entryEntrySeq.SetRDFI("231380104")
			entryEntrySeq.DFIAccountNumber = "81967038518"
			entryEntrySeq.Amount = 100000
			//entryEntrySeq.IndividualName = randomdata.FullName(randomdata.RandomGender)
			entryEntrySeq.IndividualName = "Steven Tander"
			entryEntrySeq.SetTraceNumber(bh.ODFIIdentification, entrySeq)
			//entryEntrySeq.IdentificationNumber = "#" + randomdata.RandStringRunes(13) + "#"
			entryEntrySeq.IdentificationNumber = "#83738AB#"
			entryEntrySeq.Category = ach.CategoryForward

			// Add addenda record for an entry
			addendaEntrySeq := ach.NewAddenda05()
			addendaEntrySeq.PaymentRelatedInformation = "bonus pay for amazing work on #OSS"
			entryEntrySeq.AddAddenda(addendaEntrySeq)

			// Add entries
			batch.AddEntry(entryEntrySeq)

		}

		// Create the batch.
		if err := batch.Create(); err != nil {
			fmt.Printf("%T: %s", err, err)
		}

		// Add batch to the file
		file.AddBatch(batch)
	}

	// ensure we have a validated file structure
	if file.Validate(); err != nil {
		fmt.Printf("Could not validate entire file: %v", err)
	}

	// Create the file
	if err := file.Create(); err != nil {
		fmt.Printf("%T: %s", err, err)
	}

	// Write to a file
	w := ach.NewWriter(f)
	if err := w.Write(file); err != nil {
		fmt.Printf("%T: %s", err, err)
	}
	w.Flush()
	if err := f.Close(); err != nil {
		fmt.Println(err.Error())
	}
}
<|endoftext|>


<commit_before>package tests

import (
	"database/sql"
	"log"
	"os"
	"reflect"
	"testing"

	_ "github.com/mattn/go-sqlite3"
	"github.com/monochromegane/argen"
)

func TestMain(m *testing.M) {
	db, err := sql.Open("sqlite3", ":memory:")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	Use(db)
	sqlStmt := "create table users (id integer not null primary key, name text);"
	_, err = db.Exec(sqlStmt)
	if err != nil {
		log.Fatal(err, sqlStmt)
	}

	os.Exit(m.Run())
}

func TestSelect(t *testing.T) {
	u := &User{Name: "test"}
	u.Save()
	defer User{}.DeleteAll()

	u, err := User{}.Select("id").First()
	assertError(t, err)

	if !ar.IsZero(u.Name) {
		t.Errorf("column value should be empty, but %s", u.Name)
	}
}

func TestFind(t *testing.T) {
	expect := &User{Name: "test"}
	expect.Save()
	defer User{}.DeleteAll()

	u, err := User{}.Find(1)
	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func TestFindBy(t *testing.T) {
	expect := &User{Name: "test"}
	expect.Save()
	defer User{}.DeleteAll()

	u, err := User{}.FindBy("name", "test")
	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func TestFirst(t *testing.T) {
	for _, name := range []string{"test1", "test2"} {
		u := &User{Name: name}
		u.Save()
	}
	defer User{}.DeleteAll()

	expect, _ := User{}.Where("name", "test1").QueryRow()

	u, err := User{}.First()
	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func TestLast(t *testing.T) {
	for _, name := range []string{"test1", "test2"} {
		u := &User{Name: name}
		u.Save()
	}
	defer User{}.DeleteAll()

	expect, _ := User{}.Where("name", "test2").QueryRow()

	u, err := User{}.Last()
	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func assertEqualStruct(t *testing.T, expect, actual interface{}) {
	if !reflect.DeepEqual(expect, actual) {
		t.Errorf("struct should be equal to %v, but %v", expect, actual)
	}
}

func assertError(t *testing.T, err error) {
	if err != nil {
		t.Errorf("error should be nil, but %v", err)
	}
}
<commit_msg>Added where test.<commit_after>package tests

import (
	"database/sql"
	"log"
	"os"
	"reflect"
	"testing"

	_ "github.com/mattn/go-sqlite3"
	"github.com/monochromegane/argen"
)

func TestMain(m *testing.M) {
	db, err := sql.Open("sqlite3", ":memory:")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	Use(db)
	sqlStmt := "create table users (id integer not null primary key, name text);"
	_, err = db.Exec(sqlStmt)
	if err != nil {
		log.Fatal(err, sqlStmt)
	}

	os.Exit(m.Run())
}

func TestSelect(t *testing.T) {
	u := &User{Name: "test"}
	u.Save()
	defer User{}.DeleteAll()

	u, err := User{}.Select("id").First()
	assertError(t, err)

	if !ar.IsZero(u.Name) {
		t.Errorf("column value should be empty, but %s", u.Name)
	}
}

func TestFind(t *testing.T) {
	expect := &User{Name: "test"}
	expect.Save()
	defer User{}.DeleteAll()

	u, err := User{}.Find(1)
	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func TestFindBy(t *testing.T) {
	expect := &User{Name: "test"}
	expect.Save()
	defer User{}.DeleteAll()

	u, err := User{}.FindBy("name", "test")
	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func TestFirst(t *testing.T) {
	for _, name := range []string{"test1", "test2"} {
		u := &User{Name: name}
		u.Save()
	}
	defer User{}.DeleteAll()

	expect, _ := User{}.Where("name", "test1").QueryRow()

	u, err := User{}.First()
	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func TestLast(t *testing.T) {
	for _, name := range []string{"test1", "test2"} {
		u := &User{Name: name}
		u.Save()
	}
	defer User{}.DeleteAll()

	expect, _ := User{}.Where("name", "test2").QueryRow()

	u, err := User{}.Last()
	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func TestWhere(t *testing.T) {
	expect := &User{Name: "test"}
	expect.Save()
	defer User{}.DeleteAll()

	u, err := User{}.Where("name", "test").And("id", 1).QueryRow()

	assertError(t, err)
	assertEqualStruct(t, expect, u)
}

func assertEqualStruct(t *testing.T, expect, actual interface{}) {
	if !reflect.DeepEqual(expect, actual) {
		t.Errorf("struct should be equal to %v, but %v", expect, actual)
	}
}

func assertError(t *testing.T, err error) {
	if err != nil {
		t.Errorf("error should be nil, but %v", err)
	}
}
<|endoftext|>