go-hep.org/x/hep@v0.38.1/fwk/rio/input.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  // InputStreamer reads data from a (set of) rio-stream(s)
    18  type InputStreamer struct {
    19  	Names []string            // input filenames
    20  	r     io.ReadCloser       // underlying input file(s)
    21  	rio   *rio.Reader         // input rio-stream
    22  	scan  *rio.Scanner        // input records-scanner
    23  	ports map[string]fwk.Port // input ports to read/populate
    24  }
    25  
    26  func (input *InputStreamer) Connect(ports []fwk.Port) error {
    27  	var err error
    28  
    29  	input.ports = make(map[string]fwk.Port, len(ports))
    30  
    31  	// FIXME(sbinet): handle multi-reader
    32  	// FIXME(sbinet): handle local/remote files, protocols
    33  	input.r, err = os.Open(input.Names[0])
    34  	if err != nil {
    35  		return err
    36  	}
    37  
    38  	input.rio, err = rio.NewReader(input.r)
    39  	if err != nil {
    40  		return err
    41  	}
    42  
    43  	recnames := make([]rio.Selector, 0, len(input.ports))
    44  	for _, port := range ports {
    45  		input.ports[port.Name] = port
    46  		rec := input.rio.Record(port.Name)
    47  		err = rec.Connect(port.Name, reflect.New(port.Type))
    48  		if err != nil {
    49  			return err
    50  		}
    51  		recnames = append(recnames, rio.Selector{Name: port.Name, Unpack: true})
    52  	}
    53  
    54  	input.scan = rio.NewScanner(input.rio)
    55  	input.scan.Select(recnames)
    56  	return err
    57  }
    58  
    59  func (input *InputStreamer) Read(ctx fwk.Context) error {
    60  	store := ctx.Store()
    61  	recs := make(map[string]struct{}, len(input.ports))
    62  	for range len(input.ports) {
    63  		if !input.scan.Scan() {
    64  			err := input.scan.Err()
    65  			if err == nil {
    66  				return io.EOF
    67  			}
    68  		}
    69  		rec := input.scan.Record()
    70  		blk := rec.Block(rec.Name())
    71  		obj := reflect.New(input.ports[rec.Name()].Type).Elem()
    72  		err := blk.Read(obj.Addr().Interface())
    73  		if err != nil {
    74  			return fmt.Errorf("block-read error: %w", err)
    75  		}
    76  		err = store.Put(rec.Name(), obj.Interface())
    77  		if err != nil {
    78  			return fmt.Errorf("store-put error: %w", err)
    79  		}
    80  		recs[rec.Name()] = struct{}{}
    81  	}
    82  
    83  	if len(recs) != len(input.ports) {
    84  		return fmt.Errorf("fwk.rio: expected inputs: %d, got: %d", len(input.ports), len(recs))
    85  	}
    86  
    87  	return nil
    88  }
    89  
    90  func (input *InputStreamer) Disconnect() error {
    91  	var err error
    92  	// make sure we don't leak filedescriptors
    93  	defer input.r.Close()
    94  
    95  	err = input.rio.Close()
    96  	if err != nil {
    97  		return err
    98  	}
    99  
   100  	err = input.r.Close()
   101  	if err != nil {
   102  		return err
   103  	}
   104  
   105  	return err
   106  }