go-hep.org/x/hep@v0.38.1/fwk/rio/output.go (about)

     1  // Copyright ©2015 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 rio
     6  
     7  import (
     8  	"fmt"
     9  	"io"
    10  	"os"
    11  	"reflect"
    12  
    13  	"go-hep.org/x/hep/fwk"
    14  	"go-hep.org/x/hep/rio"
    15  )
    16  
    17  // OutputStreamer writes data to a rio-stream.
    18  type OutputStreamer struct {
    19  	Name  string         // output filename
    20  	w     io.WriteCloser // underlying output file
    21  	rio   *rio.Writer    // output rio-stream
    22  	recs  []*rio.Record  // list of connected records to write out
    23  	ports []fwk.Port
    24  }
    25  
    26  func (o *OutputStreamer) Connect(ports []fwk.Port) error {
    27  	var err error
    28  
    29  	o.ports = make([]fwk.Port, len(ports))
    30  	copy(o.ports, ports)
    31  
    32  	// FIXME(sbinet): handle local/remote files, protocols
    33  	o.w, err = os.Create(o.Name)
    34  	if err != nil {
    35  		return err
    36  	}
    37  
    38  	o.rio, err = rio.NewWriter(o.w)
    39  	if err != nil {
    40  		return err
    41  	}
    42  
    43  	for _, port := range o.ports {
    44  		rec := o.rio.Record(port.Name)
    45  		err = rec.Connect(port.Name, reflect.New(port.Type))
    46  		if err != nil {
    47  			return err
    48  		}
    49  		o.recs = append(o.recs, rec)
    50  	}
    51  
    52  	return err
    53  }
    54  
    55  func (o *OutputStreamer) Disconnect() error {
    56  	// make sure we don't leak filedescriptors
    57  	defer o.w.Close()
    58  
    59  	err := o.rio.Close()
    60  	if err != nil {
    61  		return err
    62  	}
    63  
    64  	err = o.w.Close()
    65  	if err != nil {
    66  		return err
    67  	}
    68  
    69  	return err
    70  }
    71  
    72  func (o *OutputStreamer) Write(ctx fwk.Context) error {
    73  	var err error
    74  	store := ctx.Store()
    75  
    76  	for i, rec := range o.recs {
    77  		port := o.ports[i]
    78  
    79  		n := rec.Name()
    80  		blk := rec.Block(n)
    81  		obj, err := store.Get(n)
    82  		if err != nil {
    83  			return err
    84  		}
    85  
    86  		rt := reflect.TypeOf(obj)
    87  		if rt != port.Type {
    88  			return fmt.Errorf("record[%s]: got type=%q, want type=%q",
    89  				rec.Name(),
    90  				rt.Name(),
    91  				port.Type,
    92  			)
    93  		}
    94  
    95  		err = blk.Write(obj)
    96  		if err != nil {
    97  			return err
    98  		}
    99  
   100  		err = rec.Write()
   101  		if err != nil {
   102  			return err
   103  		}
   104  	}
   105  	return err
   106  }