csv2root code for root-forum/49692

main.go

// code for https://root-forum.cern.ch/t/makecsvdataframe-to-convert-from-csv-to-root-file/49692
package main

import (
	"bufio"
	"encoding/csv"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"

	"go-hep.org/x/hep/groot"
	"go-hep.org/x/hep/groot/rtree"
)

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

	var (
		oname = flag.String("o", "output.root", "path to output ROOT file")
		tname = flag.String("t", "irsn", "name of the output Tree")
	)

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Usage: csv2root [OPTIONS] INPUT.csv

ex:
 $> csv2root -o out.root -t mytree ./data.csv

options:
`)
		flag.PrintDefaults()
	}

	flag.Parse()

	run(*oname, *tname, flag.Arg(0))
}

func run(oname, tname, fname string) {
	f, err := os.Open(fname)
	if err != nil {
		log.Fatalf("could not open input CSV file: %+v", err)
	}
	defer f.Close()

	r := csv.NewReader(bufio.NewReader(f))
	r.Comma = ';'

	hdr, err := r.Read()
	if err != nil {
		log.Fatalf("could not read CSV header: %+v", err)
	}
	log.Printf("read CSV header: %q", hdr)

	o, err := groot.Create(oname)
	if err != nil {
		log.Fatalf("could not create output ROOT file: %+v", err)
	}
	defer o.Close()

	var (
		evt   Evt
		wvars = rtree.WriteVarsFromStruct(&evt)
	)

	tree, err := rtree.NewWriter(o, tname, wvars)
	if err != nil {
		log.Fatalf("could not create tree writer: %+v", err)
	}
	defer tree.Close()

loop:
	for i := 1; true; i++ {
		if i%1000 == 0 {
			log.Printf("event %d", i)
		}
		recs, err := r.Read()
		if err != nil {
			if errors.Is(err, io.EOF) {
				log.Printf("handled %d events", i-1)
				break loop
			}
			log.Fatalf("could not read line %d from CSV: %+v", i, err)
		}
		err = evt.UnmarshalCSV(recs)
		if err != nil {
			log.Fatalf("could not decode line %d from CSV: %+v", i, err)
		}
		//	log.Printf(">>> %+v", evt)
		_, err = tree.Write()
		if err != nil {
			log.Fatalf("could not write event #%d: %+v", i, err)
		}
	}

	err = tree.Close()
	if err != nil {
		log.Fatalf("could not write output ROOT file: %+v", err)
	}
	err = o.Close()
	if err != nil {
		log.Fatalf("could not write output ROOT file: %+v", err)
	}
}

type Evt struct {
	Board    uint32
	Channel  uint32
	Timetag  uint64
	Energy   uint32
	EneShort uint32
	Flag     uint32
}

func (evt *Evt) UnmarshalCSV(recs []string) error {
	var err error
	evt.Board, err = toU32(recs[0])
	if err != nil {
		return err
	}
	evt.Channel, err = toU32(recs[1])
	if err != nil {
		return err
	}
	evt.Timetag, err = toU64(recs[1])
	if err != nil {
		return err
	}
	evt.Energy, err = toU32(recs[3])
	if err != nil {
		return err
	}
	evt.EneShort, err = toU32(recs[4])
	if err != nil {
		return err
	}
	evt.Flag, err = toU32x(recs[5])
	if err != nil {
		return err
	}

	return nil
}

func toU32(v string) (uint32, error) {
	o, err := strconv.ParseUint(v, 10, 32)
	return uint32(o), err
}

func toU32x(v string) (uint32, error) {
	o, err := strconv.ParseUint(strings.TrimLeft(v, "0x"), 16, 32)
	return uint32(o), err
}

func toU64(v string) (uint64, error) {
	return strconv.ParseUint(v, 10, 64)
}