go-hep.org/x/hep@v0.38.1/groot/cmd/root-gen-type/main.go (about)

     1  // Copyright ©2019 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:generate go run ./gen-data.go -f ../../testdata/streamers.root
     6  //go:generate go run . -p main -t Event,P3 -o ./testdata/streamers.txt ../../testdata/streamers.root
     7  
     8  // Command root-gen-type generates a Go type from the StreamerInfo contained
     9  // in a ROOT file.
    10  package main // import "go-hep.org/x/hep/groot/cmd/root-gen-type"
    11  
    12  import (
    13  	"flag"
    14  	"fmt"
    15  	"io"
    16  	"log"
    17  	"os"
    18  	"regexp"
    19  	"strings"
    20  
    21  	"go-hep.org/x/hep/groot"
    22  	"go-hep.org/x/hep/groot/rdict"
    23  	_ "go-hep.org/x/hep/groot/ztypes"
    24  )
    25  
    26  func main() {
    27  	log.SetPrefix("root-gen-type: ")
    28  	log.SetFlags(0)
    29  
    30  	var (
    31  		typeNames = flag.String("t", ".*", "comma-separated list of (regexp) type names")
    32  		pkgPath   = flag.String("p", "", "package import path")
    33  		output    = flag.String("o", "", "output file name")
    34  		verbose   = flag.Bool("v", false, "enable verbose mode")
    35  	)
    36  
    37  	flag.Usage = func() {
    38  		fmt.Fprintf(
    39  			os.Stderr,
    40  			`Usage: root-gen-type [options] input.root
    41  
    42  ex:
    43   $> root-gen-type -p mypkg -t MyType -o streamers_gen.go ./input.root
    44  
    45  options:
    46  `,
    47  		)
    48  		flag.PrintDefaults()
    49  	}
    50  
    51  	flag.Parse()
    52  
    53  	if *typeNames == "" || *pkgPath == "" {
    54  		flag.Usage()
    55  		os.Exit(2)
    56  	}
    57  
    58  	if flag.NArg() != 1 {
    59  		flag.Usage()
    60  		os.Exit(2)
    61  	}
    62  
    63  	types := strings.Split(*typeNames, ",")
    64  
    65  	var (
    66  		err error
    67  		out io.WriteCloser
    68  	)
    69  
    70  	switch *output {
    71  	case "":
    72  		out = os.Stdout
    73  	default:
    74  		out, err = os.Create(*output)
    75  		if err != nil {
    76  			log.Fatal(err)
    77  		}
    78  		defer out.Close()
    79  	}
    80  
    81  	err = generate(out, *pkgPath, types, flag.Arg(0), *verbose)
    82  	if err != nil {
    83  		log.Fatal(err)
    84  	}
    85  
    86  	err = out.Close()
    87  	if err != nil {
    88  		log.Fatal(err)
    89  	}
    90  }
    91  
    92  func generate(w io.Writer, pkg string, types []string, fname string, verbose bool) error {
    93  	f, err := groot.Open(fname)
    94  	if err != nil {
    95  		return err
    96  	}
    97  
    98  	g, err := rdict.NewGenGoType(pkg, f, verbose)
    99  	if err != nil {
   100  		return err
   101  	}
   102  
   103  	filters := make([]*regexp.Regexp, len(types))
   104  	for i, t := range types {
   105  		filters[i] = regexp.MustCompile(t)
   106  	}
   107  
   108  	accept := func(name string) string {
   109  		for _, filter := range filters {
   110  			if filter.MatchString(name) {
   111  				return name
   112  			}
   113  		}
   114  		return ""
   115  	}
   116  	for _, si := range f.StreamerInfos() {
   117  		if t := accept(si.Name()); t != "" {
   118  			err := g.Generate(t)
   119  			if err != nil {
   120  				log.Fatal(err)
   121  			}
   122  		}
   123  	}
   124  
   125  	buf, err := g.Format()
   126  	if err != nil {
   127  		return err
   128  	}
   129  
   130  	_, err = w.Write(buf)
   131  	if err != nil {
   132  		return err
   133  	}
   134  
   135  	return nil
   136  }