go-hep.org/x/hep@v0.38.1/hepmc/rootcnv/writer.go (about)

     1  // Copyright ©2022 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  package rootcnv
     6  
     7  import (
     8  	"fmt"
     9  
    10  	"go-hep.org/x/hep/groot/riofs"
    11  	"go-hep.org/x/hep/groot/rtree"
    12  	"go-hep.org/x/hep/hepmc"
    13  )
    14  
    15  // FlatTreeWriter writes HepMC events as a flat ROOT TTree.
    16  type FlatTreeWriter struct {
    17  	w rtree.Writer
    18  
    19  	evt   event
    20  	wvars []rtree.WriteVar
    21  }
    22  
    23  // NewFlatTreeWriter creates a new named tree under the dir directory.
    24  func NewFlatTreeWriter(dir riofs.Directory, name string, opts ...rtree.WriteOption) (*FlatTreeWriter, error) {
    25  	var w FlatTreeWriter
    26  	w.wvars = rtree.WriteVarsFromStruct(&w.evt)
    27  	tree, err := rtree.NewWriter(dir, name, w.wvars, opts...)
    28  	if err != nil {
    29  		return nil, fmt.Errorf("hepmc: could not create flat-tree writer %q: %w", name, err)
    30  	}
    31  
    32  	w.w = tree
    33  
    34  	return &w, nil
    35  }
    36  
    37  func (w *FlatTreeWriter) Close() error {
    38  	w.wvars = nil
    39  	return w.w.Close()
    40  }
    41  
    42  func (w *FlatTreeWriter) Write(evt hepmc.Event) error {
    43  	w.evt.reset()
    44  
    45  	err := w.evt.read(&evt)
    46  	if err != nil {
    47  		return fmt.Errorf("hepmc: could not encode event to ROOT: %w", err)
    48  	}
    49  
    50  	_, err = w.w.Write()
    51  	if err != nil {
    52  		return fmt.Errorf("hepmc: could not write event to ROOT: %w", err)
    53  	}
    54  
    55  	return nil
    56  }
    57  
    58  var (
    59  	_ hepmc.Writer = (*FlatTreeWriter)(nil)
    60  )